GH-50432: [Ruby] Add ArrowFormat::{Array,Bitmap}#==

Comparing 2 bitmaps/arrays is a common operation. We should provide
content based (not identity based) comparition feature.

Author: kou

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

@@ -24,6 +24,8 @@ module ArrowFormat
   using FlatBuffers::AppendAsBytes if FlatBuffers.const_defined?(:AppendAsBytes)
 
   class Array
+    include Enumerable
+
     attr_reader :type
     attr_reader :size
     alias_method :length, :size
@@ -64,18 +66,29 @@ def empty?
       @size.zero?
     end
 
+    def ==(other)
+      return false unless other.is_a?(self.class)
+      return true if @validity_buffer.nil? and other.validity_buffer.nil?
+      if @offset == other.offset and
+        @size == other.size and
+        @validity_buffer == other.validity_buffer
+        return true
+      end
+      validity_bitmap == other.validity_bitmap
+    end
+
     protected
     def slice!(offset, size)
       @offset = offset
       @size = size
       clear_cache
     end
 
-    private
     def validity_bitmap
       @validity_bitmap ||= Bitmap.new(@validity_buffer, @offset, @size)
     end
 
+    private
     def apply_validity(array)
       return array if @validity_buffer.nil?
       validity_bitmap.each_with_index do |is_valid, i|
@@ -159,11 +172,20 @@ def n_nulls
     def to_a
       [nil] * @size
     end
+
+    def each
+      return to_enum(__method__) unless block_given?
+
+      @size.times do
+        yield(nil)
+      end
+    end
   end
 
   class PrimitiveArray < Array
     include BufferAlignable
 
+    attr_reader :values_buffer
     def initialize(*args)
       n_args = args.size
       if self.class.respond_to?(:type)
@@ -173,7 +195,7 @@ def initialize(*args)
         type = args.shift
         expected_n_args = "2 or 4"
       end
-      args = build_data(args[0], type) if args.size == 1
+      args = build_data(args[0], type) if n_args == 1
       if args.size != 3
         message =
           "wrong number of arguments " +
@@ -192,15 +214,49 @@ def to_a
       apply_validity(@values_buffer.values(@type.buffer_type, offset, @size))
     end
 
+    def each(&block)
+      return to_enum(__method__) {@size} unless block_given?
+
+      each_value = Enumerator.new(@size) do |yielder|
+        offset = element_size * @offset
+        @values_buffer.each(@type.buffer_type, offset, @size) do |_, value|
+          yielder << value
+        end
+      end
+      if @validity_buffer.nil?
+        each_value.each(&block)
+      else
+        validity_bitmap.zip(each_value) do |is_valid, value|
+          if is_valid
+            yield(value)
+          else
+            yield(nil)
+          end
+        end
+      end
+    end
+
     def each_buffer
-      return to_enum(__method__) unless block_given?
+      return to_enum(__method__) {2} unless block_given?
 
       yield(slice_bitmap_buffer(:validity, @validity_buffer))
       yield(slice_fixed_element_size_buffer(:values,
                                             @values_buffer,
                                             element_size))
     end
 
+    def ==(other)
+      return false unless super(other)
+      if @offset == other.offset and
+        @size == other.size and
+        @values_buffer == other.values_buffer
+        return true
+      end
+      zip(other).all? do |value, other_value|
+        value == other_value
+      end
+    end
+
     private
     def element_size
       IO::Buffer.size_of(@type.buffer_type)
@@ -243,8 +299,26 @@ def to_a
       apply_validity(values)
     end
 
+
+    def each(&block)
+      return to_enum(__method__) {@size} unless block_given?
+
+      @values_bitmap ||= Bitmap.new(@values_buffer, @offset, @size)
+      if @validity_buffer.nil?
+        @values_bitmap.each(&block)
+      else
+        validity_bitmap.zip(@values_bitmap) do |is_valid, value|
+          if is_valid
+            yield(value)
+          else
+            yield(nil)
+          end
+        end
+      end
+    end
+
     def each_buffer
-      return to_enum(__method__) unless block_given?
+      return to_enum(__method__) {2} unless block_given?
 
       yield(slice_bitmap_buffer(:validity, @validity_buffer))
       yield(slice_bitmap_buffer(:values, @values_buffer))

ruby/red-arrow-format/lib/arrow-format/bitmap.rb

@@ -19,10 +19,21 @@ module ArrowFormat
   class Bitmap
     include Enumerable
 
-    def initialize(buffer, offset, n_values)
+    attr_reader :offset
+    attr_reader :size
+    alias_method :length, :size
+    def initialize(buffer, offset, size)
       @buffer = buffer
       @offset = offset
-      @n_values = n_values
+      @size = size
+    end
+
+    def ==(other)
+      return false unless other.is_a?(self.class)
+      return false unless @size == other.size
+      zip(other).all? do |bit, other_bit|
+        bit == other_bit
+      end
     end
 
     def [](i)
@@ -31,19 +42,19 @@ def [](i)
     end
 
     def each
-      return to_enum(__method__) unless block_given?
+      return to_enum(__method__) {@size} unless block_given?
 
       # TODO: Optimize
       current = -1
-      n_bytes = (@offset + @n_values) / 8
+      n_bytes = (@offset + @size) / 8
       @buffer.each(:U8, 0, n_bytes) do |offset, value|
         8.times do |i|
           current += 1
           next if current < @offset
           yield((value & (1 << (i % 8))) > 0)
         end
       end
-      remained_bits = (@offset + @n_values) % 8
+      remained_bits = (@offset + @size) % 8
       unless remained_bits.zero?
         value = @buffer.get_value(:U8, n_bytes)
         remained_bits.times do |i|

ruby/red-arrow-format/test/test-bitmap.rb

@@ -0,0 +1,34 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+class TestBitmap < Test::Unit::TestCase
+  sub_test_case("#==") do
+    def test_no_slice
+      buffer1 = IO::Buffer.for("\xf0")
+      buffer2 = IO::Buffer.for("\xf0")
+      assert_equal(ArrowFormat::Bitmap.new(buffer1, 0, 8),
+                   ArrowFormat::Bitmap.new(buffer2, 0, 8))
+    end
+
+    def test_sliced
+      buffer1 = IO::Buffer.for("\xf0")
+      buffer2 = IO::Buffer.for("\x00\xf0")
+      assert_equal(ArrowFormat::Bitmap.new(buffer1, 0, 8),
+                   ArrowFormat::Bitmap.new(buffer2, 8, 8))
+    end
+  end
+end

ruby/red-arrow-format/test/test-boolean-array.rb

@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+class TestBooleanArray < Test::Unit::TestCase
+  sub_test_case("#initialize") do
+    def test_no_null
+      assert_equal([true, false],
+                   ArrowFormat::BooleanArray.new([true, false]).to_a)
+    end
+
+    def test_null_multiple_bytes
+      values = [true] * 8 + [nil, false]
+      assert_equal(values,
+                   ArrowFormat::BooleanArray.new(values).to_a)
+    end
+
+    def test_mixed
+      assert_equal([true, nil, false],
+                   ArrowFormat::BooleanArray.new([true, nil, false]).to_a)
+    end
+  end
+
+  sub_test_case("#==") do
+    def test_no_slice
+      array1 = ArrowFormat::BooleanArray.new([true, nil, false])
+      array2 = ArrowFormat::BooleanArray.new([true, nil, false])
+      assert_equal(array1, array2)
+    end
+
+    def test_sliced
+      array1 = ArrowFormat::BooleanArray.new([true, nil, false])
+      array2 = ArrowFormat::BooleanArray.new([true, true, nil, false, true])
+      assert_equal(array1, array2.slice(1, 3))
+    end
+  end
+end

ruby/red-arrow-format/test/test-float32-array.rb

@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+class TestFloat32Array < Test::Unit::TestCase
+  sub_test_case("#initialize") do
+    def test_no_null
+      values = [-Float::INFINITY, -0.0, +0.0, +Float::INFINITY]
+      assert_equal(values,
+                   ArrowFormat::Float32Array.new(values).to_a)
+    end
+
+    def test_mixed
+      values = [-Float::INFINITY, -0.0, nil, +0.0, +Float::INFINITY]
+      assert_equal(values,
+                   ArrowFormat::Float32Array.new(values).to_a)
+    end
+  end
+
+  sub_test_case("#==") do
+    def test_no_slice
+      values = [-Float::INFINITY, -0.0, nil, +0.0, +Float::INFINITY]
+      array1 = ArrowFormat::Float32Array.new(values)
+      array2 = ArrowFormat::Float32Array.new(values)
+      assert_equal(array1, array2)
+    end
+
+    def test_sliced
+      values = [-Float::INFINITY, -0.0, nil, +0.0, +Float::INFINITY]
+      array1 = ArrowFormat::Float32Array.new(values)
+      array2 = ArrowFormat::Float32Array.new([0.0] + values + [0.0])
+      assert_equal(array1, array2.slice(1, 5))
+    end
+  end
+end

ruby/red-arrow-format/test/test-float64-array.rb

@@ -0,0 +1,48 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+class TestFloat64Array < Test::Unit::TestCase
+  sub_test_case("#initialize") do
+    def test_no_null
+      values = [-Float::INFINITY, -0.0, +0.0, +Float::INFINITY]
+      assert_equal(values,
+                   ArrowFormat::Float64Array.new(values).to_a)
+    end
+
+    def test_mixed
+      values = [-Float::INFINITY, -0.0, nil, +0.0, +Float::INFINITY]
+      assert_equal(values,
+                   ArrowFormat::Float64Array.new(values).to_a)
+    end
+  end
+
+  sub_test_case("#==") do
+    def test_no_slice
+      values = [-Float::INFINITY, -0.0, nil, +0.0, +Float::INFINITY]
+      array1 = ArrowFormat::Float64Array.new(values)
+      array2 = ArrowFormat::Float64Array.new(values)
+      assert_equal(array1, array2)
+    end
+
+    def test_sliced
+      values = [-Float::INFINITY, -0.0, nil, +0.0, +Float::INFINITY]
+      array1 = ArrowFormat::Float64Array.new(values)
+      array2 = ArrowFormat::Float64Array.new([0.0] + values + [0.0])
+      assert_equal(array1, array2.slice(1, 5))
+    end
+  end
+end