red-arrow-format 25.0.0 → 26.0.0

data/lib/arrow-format/bitmap-builder.rb

@@ -0,0 +1,89 @@
+# 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
+  using FlatBuffers::AppendAsBytes if FlatBuffers.const_defined?(:AppendAsBytes)
+
+  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
+      pad!(@buffer, buffer_padding_size(@buffer))
+      @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

data/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
@@ -18,23 +19,43 @@ 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
+      lazy.zip(other).all? do |bit, other_bit|
+        bit == other_bit
+      end
     end
 
     def [](i)
+      return false if @buffer.nil?
+
       i += @offset
       (@buffer.get_value(:U8, i / 8) & (1 << (i % 8))) > 0
     end
 
     def each
-      return to_enum(__method__) unless block_given?
+      return to_enum(__method__) {@size} unless block_given?
+
+      if @buffer.nil?
+        @size.times do
+          yield(false)
+        end
+        return
+      end
 
       # 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
@@ -42,7 +63,7 @@ module ArrowFormat
           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|

data/lib/arrow-format/error.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

data/lib/arrow-format/field.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
@@ -23,7 +24,8 @@ module ArrowFormat
                    type,
                    nullable: true,
                    metadata: nil)
-      @name = name
+      name = name.to_s if name.is_a?(Symbol)
+      @name = name.to_str
       @type = type
       @nullable = nullable
       @metadata = metadata

data/lib/arrow-format/record-batch.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
@@ -15,6 +16,7 @@
 # under the License.
 
 require_relative "buffer-alignable"
+require_relative "record"
 
 module ArrowFormat
   class RecordBatch
@@ -26,7 +28,19 @@ module ArrowFormat
     alias_method :length, :n_rows
     attr_reader :columns
     attr_reader :message_metadata
-    def initialize(schema, n_rows, columns, message_metadata: nil)
+    def initialize(*args, message_metadata: nil)
+      n_args = args.size
+      case n_args
+      when 1
+        args = build(args[0])
+      when 2
+        args = build_with_schema(*args)
+      end
+      if args.size != 3
+        message = "wrong number of arguments (given #{n_args}, expected 1..3)"
+        raise ArgumentError, message
+      end
+      schema, n_rows, columns = args
       @schema = schema
       @n_rows = n_rows
       @columns = columns
@@ -37,12 +51,43 @@ module ArrowFormat
       @n_rows.zero?
     end
 
-    def to_h
-      hash = {}
-      @schema.fields.zip(@columns) do |field, column|
-        hash[field.name] = column
+    def find_column(name_or_index)
+      case name_or_index
+      when Integer
+        @columns[name_or_index]
+      when Symbol
+        name_to_column[name_or_index.to_s]
+      else
+        name_to_column[name_or_index.to_str]
+      end
+    end
+
+    def each_record(reuse_record: false)
+      unless block_given?
+        return to_enum(__method__, reuse_record: reuse_record)
+      end
+
+      if reuse_record
+        record = Record.new(self, nil)
+        @n_rows.times do |i|
+          record.index = i
+          yield(record)
+        end
+      else
+        @n_rows.times do |i|
+          yield(Record.new(self, i))
+        end
+      end
+    end
+
+    def records
+      size.times.collect do |i|
+        Record.new(self, i)
       end
-      hash
+    end
+
+    def to_h
+      name_to_column.dup
     end
 
     def to_flatbuffers
@@ -101,5 +146,121 @@ module ArrowFormat
         end
       end
     end
+
+    private
+    def build_with_schema(schema, data)
+      fields = schema.fields
+      name_to_index = {}
+      fields.each_with_index do |field, i|
+        name_to_index[field.name] = i
+      end
+      if data.is_a?(Hash)
+        raw_columns = []
+        data.each do |name, values|
+          field_index = name_to_index[name.to_s]
+          raw_columns[field_index] = values if field_index
+        end
+        columns = fields.zip(raw_columns).collect do |field, values|
+          field.type.build_array(values || [])
+        end
+        all_n_rows = columns.collect(&:size)
+        if all_n_rows.uniq.size != 1
+          message = "inconsistent the number of rows: #{all_n_rows.join(", ")}"
+          raise ArgumentError, message
+        end
+      else
+        raw_columns = fields.collect { [] }
+        data.each_with_index do |record, nth_record|
+          case record
+          when nil
+          when Hash
+            record.each do |name, value|
+              field_index = name_to_index[name.to_s]
+              raw_columns[field_index] << value if field_index
+            end
+          else
+            if record.size > raw_columns.size
+              message = "row #{nth_record} has more values than schema fields"
+              raise ArgumentError, message
+            end
+            record.each_with_index do |value, field_index|
+              raw_columns[field_index] << value
+            end
+          end
+          raw_columns.each do |column|
+            column << nil if column.size != nth_record + 1
+          end
+        end
+        columns = fields.zip(raw_columns).collect do |field, values|
+          field.type.build_array(values)
+        end
+      end
+      n_rows = columns.first&.size || 0
+      [schema, n_rows, columns]
+    end
+
+    def build(data)
+      records = nil
+      fields = []
+      columns = []
+      mode = nil
+      i = 0
+      data.each do |name, column|
+        if i.zero?
+          if column.nil?
+            mode = :record
+            records = {}
+          else
+            mode = :column
+          end
+        end
+
+        if mode == :record
+          record = name
+          record.each do |n, value|
+            values = (records[n] ||= [])
+            while values.size < i
+              values << nil
+            end
+            values << value
+          end
+        else
+          fields << Field.new(name, column.type)
+          columns << column
+        end
+        i += 1
+      end
+
+      if mode == :record
+        records.each do |name, values|
+          column = Array.build(values)
+          fields << Field.new(name, column.type)
+          columns << column
+        end
+      end
+
+      raise ArgumentError, "no data" if columns.empty?
+      all_n_rows = columns.collect(&:size)
+      if all_n_rows.uniq.size != 1
+        message =
+          "inconsistent the number of rows: #{all_n_rows.join(", ")}"
+        raise ArgumentError, message
+      end
+
+      return Schema.new(fields), all_n_rows[0], columns
+    end
+
+    def name_to_column
+      @name_to_column ||= build_name_to_column
+    end
+
+    def build_name_to_column
+      name_to_column = {}
+      @schema.fields.zip(@columns) do |field, column|
+        name_to_column[field.name] = column
+      end
+      name_to_column.freeze
+      name_to_column
+    end
   end
 end

data/lib/arrow-format/record.rb

@@ -0,0 +1,60 @@
+# 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.
+
+module ArrowFormat
+  class Record
+    attr_reader :record_batch
+    attr_accessor :index
+    def initialize(record_batch, index)
+      @record_batch = record_batch
+      @index = index
+    end
+
+    def [](name_or_index)
+      column = @record_batch.find_column(name_or_index)
+      return nil if column.nil?
+      column[@index]
+    end
+
+    def to_a
+      @record_batch.columns.collect do |column|
+        column[@index]
+      end
+    end
+
+    def to_h
+      attributes = {}
+      @record_batch.schema.fields.zip(@record_batch.columns) do |field, column|
+        attributes[field.name] = column[@index]
+      end
+      attributes
+    end
+
+    def respond_to_missing?(name, include_private)
+      return true if @record_batch.find_column(name)
+      super
+    end
+
+    def method_missing(name, *args, &block)
+      if args.empty?
+        column = @record_batch.find_column(name)
+        return column[@index] if column
+      end
+      super
+    end
+  end
+end

data/lib/arrow-format/schema.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
@@ -25,6 +26,17 @@ module ArrowFormat
       @message_metadata = message_metadata
     end
 
+    def [](name_or_index)
+      case name_or_index
+      when Integer
+        @fields[name_or_index]
+      when Symbol
+        name_to_field[name_or_index.to_s]
+      else
+        name_to_field[name_or_index.to_str]
+      end
+    end
+
     def to_flatbuffers
       fb_schema = FB::Schema::Data.new
       fb_schema.endianness = FB::Endianness::LITTLE
@@ -33,5 +45,18 @@ module ArrowFormat
       # fb_schema.features = @features
       fb_schema
     end
+
+    private
+    def name_to_field
+      @name_to_field ||= build_name_to_field
+    end
+
+    def build_name_to_field
+      name_to_field = {}
+      @fields.each do |field|
+        name_to_field[field.name] = field
+      end
+      name_to_field
+    end
   end
 end