GH-50385: [Ruby] Add bitmap builder for red-arrow-format

Author: kou

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

@@ -1,3 +1,4 @@
+# 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
@@ -17,6 +18,7 @@
 require "bigdecimal"
 
 require_relative "bitmap"
+require_relative "bitmap-builder"
 
 module ArrowFormat
   class Array
@@ -186,7 +188,16 @@ def element_size
   end
 
   class BooleanArray < PrimitiveArray
-    def initialize(size, validity_buffer, values_buffer)
+    def initialize(*args)
+      if args.size == 1
+        args = build_data(args[0])
+      end
+      n_args = args.size
+      if args.size != 3
+        message = "wrong number of arguments (given #{n_args}, expected 1 or 3)"
+        raise ArgumentError, message
+      end
+      size, validity_buffer, values_buffer = args
       super(BooleanType.singleton, size, validity_buffer, values_buffer)
     end
 
@@ -210,6 +221,26 @@ def clear_cache
       super
       @values_bitmap = nil
     end
+
+    def build_data(data)
+      n = 0
+      validity_buffer_builder = nil
+      values_buffer_builder = DenseBitmapBuilder.new
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+          values_buffer_builder.append(false)
+        elsif value
+          values_buffer_builder.append(true)
+        else
+          values_buffer_builder.append(false)
+        end
+        n += 1
+      end
+      validity_buffer = validity_buffer_builder&.finish(n)
+      return n, validity_buffer, values_buffer_builder.finish
+    end
   end
 
   class IntArray < PrimitiveArray

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

@@ -0,0 +1,88 @@
+# 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.
+
+require_relative "buffer-alignable"
+
+module ArrowFormat
+  class DenseBitmapBuilder
+    include BufferAlignable
+
+    def initialize
+      @buffer = +"".b
+      @n_bits = 0
+      @byte = 0
+    end
+
+    def append(value)
+      @byte |= 1 << @n_bits if value
+      @n_bits += 1
+      flush if @n_bits == 8
+      self
+    end
+
+    def finish
+      flush if @n_bits > 0
+      padding_size = buffer_padding_size(@buffer)
+      @buffer.append_as_bytes(padding(padding_size)) if padding_size > 0
+      @buffer.freeze
+      IO::Buffer.for(@buffer)
+    end
+
+    private
+    def flush
+      @buffer.append_as_bytes([@byte].pack("C"))
+      @n_bits = 0
+      @byte = 0
+    end
+  end
+
+  class SparseBitmapBuilder
+    include BufferAlignable
+
+    def initialize
+      @unset_indexes = {}
+    end
+
+    def unset(index)
+      @unset_indexes[index] = true
+    end
+
+    def finish(size)
+      builder = DenseBitmapBuilder.new
+      set_value = true
+      unset_value = false
+      if @unset_indexes.empty?
+        size.times do
+          builder.append(set_value)
+        end
+      else
+        previous_index = 0
+        @unset_indexes.keys.sort.each do |index|
+          previous_index.upto(index - 1) do
+            builder.append(set_value)
+          end
+          builder.append(unset_value)
+          previous_index = index + 1
+        end
+        (size - previous_index).times do
+          builder.append(set_value)
+        end
+      end
+      builder.finish
+    end
+  end
+end

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

@@ -1,3 +1,4 @@
+# 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

ruby/red-arrow-format/lib/arrow-format/integration/json-reader.rb

@@ -205,18 +205,11 @@ def read_schema
       end
 
       def read_bitmap(bitmap)
-        buffer = +"".b
-        bitmap.each_slice(8) do |bits|
-          byte = 0
-          while bits.size < 8
-            bits << 0
-          end
-          bits.reverse_each do |bit|
-            byte = (byte << 1) + bit
-          end
-          buffer << [byte].pack("C")
+        builder = DenseBitmapBuilder.new
+        bitmap.each do |bit|
+          builder.append(bit == 1 ? true : false)
         end
-        IO::Buffer.for(buffer)
+        builder.finish
       end
 
       def read_types(types)

ruby/red-arrow-format/test/test-boolean-array.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 TestBooleanArray < Test::Unit::TestCase
+  def test_mixed
+    assert_equal([true, nil, false],
+                 ArrowFormat::BooleanArray.new([true, nil, false]).to_a)
+  end
+
+  def test_no_null
+    assert_equal([true, false],
+                 ArrowFormat::BooleanArray.new([true, false]).to_a)
+  end
+
+  def test_more_8bits
+    values = [true] * 8 + [nil, false]
+    assert_equal(values,
+                 ArrowFormat::BooleanArray.new(values).to_a)
+  end
+end

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

@@ -0,0 +1,45 @@
+# 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 TestDenseBitmapBuilder < Test::Unit::TestCase
+  def setup
+    @builder = ArrowFormat::DenseBitmapBuilder.new
+  end
+
+  def test_empty
+    assert_equal(IO::Buffer.for(""), @builder.finish)
+  end
+
+  def test_1byte
+    8.times do |i|
+      @builder.append(i.odd?)
+    end
+    buffer = [0b10101010].pack("C") + "\x00" * 63
+    assert_equal(IO::Buffer.for(buffer),
+                 @builder.finish)
+  end
+
+  def test_9bits
+    8.times do |i|
+      @builder.append(i.odd?)
+    end
+    @builder.append(true)
+    buffer = [0b10101010].pack("C") + [0b00000001].pack("C") + "\x00" * 62
+    assert_equal(IO::Buffer.for(buffer),
+                 @builder.finish)
+  end
+end

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

@@ -0,0 +1,58 @@
+# 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 TestSparseBitmapBuilder < Test::Unit::TestCase
+  def setup
+    @builder = ArrowFormat::SparseBitmapBuilder.new
+  end
+
+  def test_empty
+    buffer = [0b00000001].pack("C") + "\x00" * 63
+    assert_equal(IO::Buffer.for(buffer),
+                 @builder.finish(1))
+  end
+
+  def test_unset_0bit
+    @builder.unset(0)
+    buffer = [0b11111110].pack("C") + "\x00" * 63
+    assert_equal(IO::Buffer.for(buffer),
+                 @builder.finish(8))
+  end
+
+  def test_unset_multiple_bytes
+    @builder.unset(15)
+    @builder.unset(31)
+    buffer =  [0b11111111].pack("C")
+    buffer += [0b01111111].pack("C")
+    buffer += [0b11111111].pack("C")
+    buffer += [0b01111111].pack("C")
+    buffer += [0b11111111].pack("C")
+    buffer += "\x00" * 59
+    assert_equal(IO::Buffer.for(buffer),
+                 @builder.finish(40))
+  end
+
+  def test_unset_duplicated
+    @builder.unset(15)
+    @builder.unset(15)
+    buffer =  [0b11111111].pack("C")
+    buffer += [0b01111111].pack("C")
+    buffer += "\x00" * 62
+    assert_equal(IO::Buffer.for(buffer),
+                 @builder.finish(16))
+  end
+end