Skip to content

Commit de8ba54

Browse files
koukou
committed
GH-50437: [Ruby] Add ArrowFormat::*IntervalArray.new(values) 
1 parent 58a6ea2 commit de8ba54

5 files changed

Lines changed: 346 additions & 7 deletions

File tree

‎ruby/red-arrow-format/lib/arrow-format/array.rb‎

Lines changed: 95 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -271,9 +271,9 @@ def build_data(data, type)
271271
if value.nil?
272272
validity_buffer_builder ||= SparseBitmapBuilder.new
273273
validity_buffer_builder.unset(i)
274-
buffer.append_as_bytes([0].pack(pack_template))
274+
buffer.append_as_bytes(pack_value(nil, pack_template))
275275
else
276-
buffer.append_as_bytes([value].pack(pack_template))
276+
buffer.append_as_bytes(pack_value(value, pack_template))
277277
end
278278
n += 1
279279
end
@@ -282,6 +282,11 @@ def build_data(data, type)
282282
buffer.freeze
283283
return n, validity_buffer, IO::Buffer.for(buffer)
284284
end
285+
286+
def pack_value(value, template)
287+
value = 0 if value.nil?
288+
[value].pack(template)
289+
end
285290
end
286291

287292
class BooleanArray < PrimitiveArray
@@ -491,29 +496,80 @@ class IntervalArray < TemporalArray
491496
end
492497

493498
class YearMonthIntervalArray < IntervalArray
499+
class << self
500+
def type
501+
YearMonthIntervalType.singleton
502+
end
503+
end
494504
end
495505

496506
class DayTimeIntervalArray < IntervalArray
507+
class << self
508+
def type
509+
DayTimeIntervalType.singleton
510+
end
511+
end
512+
497513
def to_a
498514
return [] if empty?
499515

500516
offset = element_size * @offset
501517
values = @values_buffer.
502518
each(@type.buffer_type, offset, @size * 2).
503519
each_slice(2).
504-
collect do |(_, day), (_, time)|
505-
[day, time]
520+
collect do |(_, day), (_, millisecond)|
521+
[day, millisecond]
506522
end
507523
apply_validity(values)
508524
end
509525

526+
def each(&block)
527+
return to_enum(__method__) {@size} unless block_given?
528+
529+
each_value = Enumerator.new(@size) do |yielder|
530+
offset = element_size * @offset
531+
@values_buffer.
532+
each(@type.buffer_type, offset, @size * 2).
533+
each_slice(2) do |(_, day), (_, millisecond)|
534+
yielder << [day, millisecond]
535+
end
536+
end
537+
if @validity_buffer.nil?
538+
each_value.each(&block)
539+
else
540+
validity_bitmap.zip(each_value) do |is_valid, value|
541+
if is_valid
542+
yield(value)
543+
else
544+
yield(nil)
545+
end
546+
end
547+
end
548+
end
549+
510550
private
511551
def element_size
512552
super * 2
513553
end
554+
555+
def pack_value(value, template)
556+
if value.nil?
557+
[0, 0].pack(template)
558+
elsif value.is_a?(Hash)
559+
[value[:day], value[:millisecond]].pack(template)
560+
else
561+
value.pack(template)
562+
end
563+
end
514564
end
515565

516566
class MonthDayNanoIntervalArray < IntervalArray
567+
class << self
568+
def type
569+
MonthDayNanoIntervalType.singleton
570+
end
571+
end
572+
517573
def to_a
518574
return [] if empty?
519575

@@ -527,10 +583,45 @@ def to_a
527583
apply_validity(values)
528584
end
529585

586+
def each(&block)
587+
return to_enum(__method__) {@size} unless block_given?
588+
589+
each_value = Enumerator.new(@size) do |yielder|
590+
buffer_types = @type.buffer_types
591+
value_size = IO::Buffer.size_of(buffer_types)
592+
base_offset = value_size * @offset
593+
@size.times do |i|
594+
offset = base_offset + value_size * i
595+
yielder << @values_buffer.get_values(buffer_types, offset)
596+
end
597+
end
598+
if @validity_buffer.nil?
599+
each_value.each(&block)
600+
else
601+
validity_bitmap.zip(each_value) do |is_valid, value|
602+
if is_valid
603+
yield(value)
604+
else
605+
yield(nil)
606+
end
607+
end
608+
end
609+
end
610+
530611
private
531612
def element_size
532613
IO::Buffer.size_of(@type.buffer_types)
533614
end
615+
616+
def pack_value(value, template)
617+
if value.nil?
618+
[0, 0, 0].pack(template)
619+
elsif value.is_a?(Hash)
620+
[value[:month], value[:day], value[:nanosecond]].pack(template)
621+
else
622+
value.pack(template)
623+
end
624+
end
534625
end
535626

536627
class DurationArray < TemporalArray

‎ruby/red-arrow-format/lib/arrow-format/type.rb‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -642,8 +642,12 @@ def buffer_type
642642
:s32
643643
end
644644

645+
def pack_template
646+
"l"
647+
end
648+
645649
def build_array(...)
646-
YearMonthIntervalArray.new(self, ...)
650+
YearMonthIntervalArray.new(...)
647651
end
648652
end
649653

@@ -660,8 +664,12 @@ def buffer_type
660664
:s32
661665
end
662666

667+
def pack_template
668+
"ll"
669+
end
670+
663671
def build_array(...)
664-
DayTimeIntervalArray.new(self, ...)
672+
DayTimeIntervalArray.new(...)
665673
end
666674
end
667675

@@ -678,8 +686,12 @@ def buffer_types
678686
@buffer_types ||= [:s32, :s32, :s64]
679687
end
680688

689+
def pack_template
690+
"llq"
691+
end
692+
681693
def build_array(...)
682-
MonthDayNanoIntervalArray.new(self, ...)
694+
MonthDayNanoIntervalArray.new(...)
683695
end
684696
end
685697

‎ruby/red-arrow-format/test/test-day-time-interval-array.rb‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
class TestDayTimeIntervalArray < Test::Unit::TestCase
19+
sub_test_case("#initialize") do
20+
def test_no_null
21+
values = [
22+
[1, 100],
23+
[3, 300],
24+
]
25+
assert_equal(values,
26+
ArrowFormat::DayTimeIntervalArray.new(values).to_a)
27+
end
28+
29+
def test_mixed
30+
values = [
31+
[1, 100],
32+
nil,
33+
[3, 300],
34+
]
35+
assert_equal(values,
36+
ArrowFormat::DayTimeIntervalArray.new(values).to_a)
37+
end
38+
39+
def test_hash
40+
values = [
41+
{day: 1, millisecond: 100},
42+
nil,
43+
{day: 3, millisecond: 300},
44+
]
45+
assert_equal([
46+
[1, 100],
47+
nil,
48+
[3, 300],
49+
],
50+
ArrowFormat::DayTimeIntervalArray.new(values).to_a)
51+
end
52+
end
53+
54+
sub_test_case("#==") do
55+
def test_no_slice
56+
values = [
57+
[1, 100],
58+
nil,
59+
[3, 300],
60+
]
61+
array1 = ArrowFormat::DayTimeIntervalArray.new(values)
62+
array2 = ArrowFormat::DayTimeIntervalArray.new(values)
63+
assert_equal(array1, array2)
64+
end
65+
66+
def test_sliced
67+
values = [
68+
[1, 100],
69+
nil,
70+
[3, 300],
71+
]
72+
array1 = ArrowFormat::DayTimeIntervalArray.new(values)
73+
array2 = ArrowFormat::DayTimeIntervalArray.new([nil, *values, nil])
74+
assert_equal(array1, array2.slice(1, 3))
75+
end
76+
77+
def test_sliced_different_content
78+
values = [
79+
[1, 100],
80+
nil,
81+
[3, 300],
82+
]
83+
array1 = ArrowFormat::DayTimeIntervalArray.new(values)
84+
array2 = ArrowFormat::DayTimeIntervalArray.new([nil, nil, *values, nil])
85+
assert_not_equal(array1, array2.slice(1, 3))
86+
end
87+
end
88+
end

‎ruby/red-arrow-format/test/test-month-day-nano-interval-array.rb‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
class TestMonthDayNanoIntervalArray < Test::Unit::TestCase
19+
sub_test_case("#initialize") do
20+
def test_no_null
21+
values = [
22+
[1, 1, 100],
23+
[3, 3, 300],
24+
]
25+
assert_equal(values,
26+
ArrowFormat::MonthDayNanoIntervalArray.new(values).to_a)
27+
end
28+
29+
def test_mixed
30+
values = [
31+
[1, 1, 100],
32+
nil,
33+
[3, 3, 300],
34+
]
35+
assert_equal(values,
36+
ArrowFormat::MonthDayNanoIntervalArray.new(values).to_a)
37+
end
38+
39+
def test_hash
40+
values = [
41+
{month: 1, day: 1, nanosecond: 100},
42+
nil,
43+
{month: 3, day: 3, nanosecond: 300},
44+
]
45+
assert_equal([
46+
[1, 1, 100],
47+
nil,
48+
[3, 3, 300],
49+
],
50+
ArrowFormat::MonthDayNanoIntervalArray.new(values).to_a)
51+
end
52+
end
53+
54+
sub_test_case("#==") do
55+
def test_no_slice
56+
values = [
57+
[1, 1, 100],
58+
nil,
59+
[3, 3, 300],
60+
]
61+
array1 = ArrowFormat::MonthDayNanoIntervalArray.new(values)
62+
array2 = ArrowFormat::MonthDayNanoIntervalArray.new(values)
63+
assert_equal(array1, array2)
64+
end
65+
66+
def test_sliced
67+
values = [
68+
[1, 1, 100],
69+
nil,
70+
[3, 3, 300],
71+
]
72+
array1 = ArrowFormat::MonthDayNanoIntervalArray.new(values)
73+
array2 = ArrowFormat::MonthDayNanoIntervalArray.new([nil, *values, nil])
74+
assert_equal(array1, array2.slice(1, 3))
75+
end
76+
77+
def test_sliced_different_content
78+
values = [
79+
[1, 1, 100],
80+
nil,
81+
[3, 3, 300],
82+
]
83+
array1 = ArrowFormat::MonthDayNanoIntervalArray.new(values)
84+
array2 = ArrowFormat::MonthDayNanoIntervalArray.new([
85+
nil,
86+
nil,
87+
*values,
88+
nil,
89+
])
90+
assert_not_equal(array1, array2.slice(1, 3))
91+
end
92+
end
93+
end

0 commit comments

Comments
 (0)