red-arrow-format 23.0.1 → 24.0.0

data/lib/arrow-format/readable.rb

@@ -17,152 +17,186 @@
 
 require_relative "array"
 require_relative "field"
+require_relative "flatbuffers"
 require_relative "record-batch"
 require_relative "schema"
 require_relative "type"
 
-require_relative "org/apache/arrow/flatbuf/binary"
-require_relative "org/apache/arrow/flatbuf/bool"
-require_relative "org/apache/arrow/flatbuf/date"
-require_relative "org/apache/arrow/flatbuf/date_unit"
-require_relative "org/apache/arrow/flatbuf/decimal"
-require_relative "org/apache/arrow/flatbuf/dictionary_encoding"
-require_relative "org/apache/arrow/flatbuf/dictionary_batch"
-require_relative "org/apache/arrow/flatbuf/duration"
-require_relative "org/apache/arrow/flatbuf/fixed_size_binary"
-require_relative "org/apache/arrow/flatbuf/floating_point"
-require_relative "org/apache/arrow/flatbuf/int"
-require_relative "org/apache/arrow/flatbuf/interval"
-require_relative "org/apache/arrow/flatbuf/interval_unit"
-require_relative "org/apache/arrow/flatbuf/large_binary"
-require_relative "org/apache/arrow/flatbuf/large_list"
-require_relative "org/apache/arrow/flatbuf/large_utf8"
-require_relative "org/apache/arrow/flatbuf/list"
-require_relative "org/apache/arrow/flatbuf/map"
-require_relative "org/apache/arrow/flatbuf/message"
-require_relative "org/apache/arrow/flatbuf/null"
-require_relative "org/apache/arrow/flatbuf/precision"
-require_relative "org/apache/arrow/flatbuf/record_batch"
-require_relative "org/apache/arrow/flatbuf/schema"
-require_relative "org/apache/arrow/flatbuf/struct_"
-require_relative "org/apache/arrow/flatbuf/time"
-require_relative "org/apache/arrow/flatbuf/time_unit"
-require_relative "org/apache/arrow/flatbuf/timestamp"
-require_relative "org/apache/arrow/flatbuf/union"
-require_relative "org/apache/arrow/flatbuf/union_mode"
-require_relative "org/apache/arrow/flatbuf/utf8"
-
 module ArrowFormat
   module Readable
     private
-    def read_schema(fb_schema)
+    def read_custom_metadata(fb_custom_metadata)
+      return nil if fb_custom_metadata.nil?
+      metadata = {}
+      fb_custom_metadata.each do |key_value|
+        metadata[key_value.key] = key_value.value
+      end
+      metadata
+    end
+
+    def read_schema(fb_schema, fb_message_custom_metadata=nil)
       fields = fb_schema.fields.collect do |fb_field|
         read_field(fb_field)
       end
-      Schema.new(fields)
+      message_metadata = read_custom_metadata(fb_message_custom_metadata)
+      Schema.new(fields,
+                 metadata: read_custom_metadata(fb_schema.custom_metadata),
+                 message_metadata: message_metadata)
     end
 
-    def read_field(fb_field)
+    def read_field(fb_field,
+                   map_entries: false,
+                   map_key: false,
+                   map_value: false)
       fb_type = fb_field.type
       case fb_type
-      when Org::Apache::Arrow::Flatbuf::Null
+      when FB::Null
         type = NullType.singleton
-      when Org::Apache::Arrow::Flatbuf::Bool
+      when FB::Bool
         type = BooleanType.singleton
-      when Org::Apache::Arrow::Flatbuf::Int
+      when FB::Int
         type = read_type_int(fb_type)
-      when Org::Apache::Arrow::Flatbuf::FloatingPoint
+      when FB::FloatingPoint
         case fb_type.precision
-        when Org::Apache::Arrow::Flatbuf::Precision::SINGLE
+        when FB::Precision::SINGLE
           type = Float32Type.singleton
-        when Org::Apache::Arrow::Flatbuf::Precision::DOUBLE
+        when FB::Precision::DOUBLE
           type = Float64Type.singleton
+        else
+          raise ReadError.new("Unsupported type: #{fb_type.inspect}")
         end
-      when Org::Apache::Arrow::Flatbuf::Date
+      when FB::Date
         case fb_type.unit
-        when Org::Apache::Arrow::Flatbuf::DateUnit::DAY
+        when FB::DateUnit::DAY
           type = Date32Type.singleton
-        when Org::Apache::Arrow::Flatbuf::DateUnit::MILLISECOND
+        when FB::DateUnit::MILLISECOND
           type = Date64Type.singleton
+        else
+          raise ReadError.new("Unsupported type: #{fb_type.inspect}")
         end
-      when Org::Apache::Arrow::Flatbuf::Time
+      when FB::Time
         case fb_type.bit_width
         when 32
           case fb_type.unit
-          when Org::Apache::Arrow::Flatbuf::TimeUnit::SECOND
+          when FB::TimeUnit::SECOND
             type = Time32Type.new(:second)
-          when Org::Apache::Arrow::Flatbuf::TimeUnit::MILLISECOND
+          when FB::TimeUnit::MILLISECOND
             type = Time32Type.new(:millisecond)
+          else
+            raise ReadError.new("Unsupported type: #{fb_type.inspect}")
           end
         when 64
           case fb_type.unit
-          when Org::Apache::Arrow::Flatbuf::TimeUnit::MICROSECOND
+          when FB::TimeUnit::MICROSECOND
             type = Time64Type.new(:microsecond)
-          when Org::Apache::Arrow::Flatbuf::TimeUnit::NANOSECOND
+          when FB::TimeUnit::NANOSECOND
             type = Time64Type.new(:nanosecond)
+          else
+            raise ReadError.new("Unsupported type: #{fb_type.inspect}")
           end
         end
-      when Org::Apache::Arrow::Flatbuf::Timestamp
+      when FB::Timestamp
         unit = fb_type.unit.name.downcase.to_sym
         type = TimestampType.new(unit, fb_type.timezone)
-      when Org::Apache::Arrow::Flatbuf::Interval
+      when FB::Interval
         case fb_type.unit
-        when Org::Apache::Arrow::Flatbuf::IntervalUnit::YEAR_MONTH
-          type = YearMonthIntervalType.new
-        when Org::Apache::Arrow::Flatbuf::IntervalUnit::DAY_TIME
-          type = DayTimeIntervalType.new
-        when Org::Apache::Arrow::Flatbuf::IntervalUnit::MONTH_DAY_NANO
-          type = MonthDayNanoIntervalType.new
+        when FB::IntervalUnit::YEAR_MONTH
+          type = YearMonthIntervalType.singleton
+        when FB::IntervalUnit::DAY_TIME
+          type = DayTimeIntervalType.singleton
+        when FB::IntervalUnit::MONTH_DAY_NANO
+          type = MonthDayNanoIntervalType.singleton
+        else
+          raise ReadError.new("Unsupported type: #{fb_type.inspect}")
         end
-      when Org::Apache::Arrow::Flatbuf::Duration
+      when FB::Duration
         unit = fb_type.unit.name.downcase.to_sym
         type = DurationType.new(unit)
-      when Org::Apache::Arrow::Flatbuf::List
+      when FB::List
         type = ListType.new(read_field(fb_field.children[0]))
-      when Org::Apache::Arrow::Flatbuf::LargeList
+      when FB::LargeList
         type = LargeListType.new(read_field(fb_field.children[0]))
-      when Org::Apache::Arrow::Flatbuf::Struct
-        children = fb_field.children.collect {|child| read_field(child)}
+      when FB::FixedSizeList
+        type = FixedSizeListType.new(read_field(fb_field.children[0]),
+                                     fb_type.list_size)
+      when FB::Struct
+        if map_entries
+          fb_children = fb_field.children
+          children = [
+            read_field(fb_children[0], map_key: true),
+            read_field(fb_children[1], map_value: true),
+          ]
+        else
+          children = fb_field.children.collect {|child| read_field(child)}
+        end
         type = StructType.new(children)
-      when Org::Apache::Arrow::Flatbuf::Union
+      when FB::Union
         children = fb_field.children.collect {|child| read_field(child)}
         type_ids = fb_type.type_ids
         case fb_type.mode
-        when Org::Apache::Arrow::Flatbuf::UnionMode::DENSE
+        when FB::UnionMode::DENSE
           type = DenseUnionType.new(children, type_ids)
-        when Org::Apache::Arrow::Flatbuf::UnionMode::SPARSE
+        when FB::UnionMode::SPARSE
           type = SparseUnionType.new(children, type_ids)
+        else
+          raise ReadError.new("Unsupported type: #{fb_type.inspect}")
         end
-      when Org::Apache::Arrow::Flatbuf::Map
-        type = MapType.new(read_field(fb_field.children[0]))
-      when Org::Apache::Arrow::Flatbuf::Binary
+      when FB::Map
+        type = MapType.new(read_field(fb_field.children[0], map_entries: true),
+                           fb_type.keys_sorted?)
+      when FB::Binary
         type = BinaryType.singleton
-      when Org::Apache::Arrow::Flatbuf::LargeBinary
+      when FB::LargeBinary
         type = LargeBinaryType.singleton
-      when Org::Apache::Arrow::Flatbuf::Utf8
+      when FB::Utf8
         type = UTF8Type.singleton
-      when Org::Apache::Arrow::Flatbuf::LargeUtf8
+      when FB::LargeUtf8
         type = LargeUTF8Type.singleton
-      when Org::Apache::Arrow::Flatbuf::FixedSizeBinary
+      when FB::FixedSizeBinary
         type = FixedSizeBinaryType.new(fb_type.byte_width)
-      when Org::Apache::Arrow::Flatbuf::Decimal
+      when FB::Decimal
         case fb_type.bit_width
         when 128
           type = Decimal128Type.new(fb_type.precision, fb_type.scale)
         when 256
           type = Decimal256Type.new(fb_type.precision, fb_type.scale)
+        else
+          raise ReadError.new("Unsupported type: #{fb_type.inspect}")
         end
+      else
+        raise ReadError.new("Unsupported type: #{fb_type.inspect}")
       end
 
       dictionary = fb_field.dictionary
       if dictionary
         dictionary_id = dictionary.id
         index_type = read_type_int(dictionary.index_type)
-        type = DictionaryType.new(index_type, type, dictionary.ordered?)
+        value_type = type
+        type = DictionaryType.new(dictionary_id,
+                                  index_type,
+                                  value_type,
+                                  dictionary.ordered?)
+      end
+
+      # Map type uses static "entries"/"key"/"value" as field names
+      # instead of field names in FlatBuffers. It's based on the
+      # specification:
+      #
+      #   The names of the child fields may be respectively "entries",
+      #   "key", and "value", but this is not enforced.
+      if map_entries
+        name = "entries"
+      elsif map_key
+        name = "key"
+      elsif map_value
+        name = "value"
       else
-        dictionary_id = nil
+        name = fb_field.name
       end
-      Field.new(fb_field.name, type, fb_field.nullable?, dictionary_id)
+      Field.new(name,
+                type,
+                nullable: fb_field.nullable?,
+                metadata: read_custom_metadata(fb_field.custom_metadata))
     end
 
     def read_type_int(fb_type)
@@ -194,17 +228,23 @@ module ArrowFormat
       end
     end
 
-    def read_record_batch(fb_record_batch, schema, body)
+    def read_record_batch(version,
+                          fb_record_batch,
+                          fb_message_custom_metadata,
+                          schema,
+                          body)
+      message_metadata = read_custom_metadata(fb_message_custom_metadata)
       n_rows = fb_record_batch.length
       nodes = fb_record_batch.nodes
       buffers = fb_record_batch.buffers
       columns = schema.fields.collect do |field|
-        read_column(field, nodes, buffers, body)
+        read_column(version, field, nodes, buffers, body)
       end
-      RecordBatch.new(schema, n_rows, columns)
+      RecordBatch.new(schema, n_rows, columns,
+                      message_metadata: message_metadata)
     end
 
-    def read_column(field, nodes, buffers, body)
+    def read_column(version, field, nodes, buffers, body)
       node = nodes.shift
       length = node.length
 
@@ -222,49 +262,64 @@ module ArrowFormat
            NumberType,
            TemporalType
         values_buffer = buffers.shift
-        values = body.slice(values_buffer.offset, values_buffer.length)
+        values = body&.slice(values_buffer.offset, values_buffer.length)
         field.type.build_array(length, validity, values)
       when VariableSizeBinaryType
         offsets_buffer = buffers.shift
         values_buffer = buffers.shift
-        offsets = body.slice(offsets_buffer.offset, offsets_buffer.length)
-        values = body.slice(values_buffer.offset, values_buffer.length)
+        offsets = body&.slice(offsets_buffer.offset, offsets_buffer.length)
+        values = body&.slice(values_buffer.offset, values_buffer.length)
         field.type.build_array(length, validity, offsets, values)
       when FixedSizeBinaryType
         values_buffer = buffers.shift
-        values = body.slice(values_buffer.offset, values_buffer.length)
+        values = body&.slice(values_buffer.offset, values_buffer.length)
         field.type.build_array(length, validity, values)
       when VariableSizeListType
         offsets_buffer = buffers.shift
-        offsets = body.slice(offsets_buffer.offset, offsets_buffer.length)
-        child = read_column(field.type.child, nodes, buffers, body)
+        offsets = body&.slice(offsets_buffer.offset, offsets_buffer.length)
+        child = read_column(version, field.type.child, nodes, buffers, body)
         field.type.build_array(length, validity, offsets, child)
+      when FixedSizeListType
+        child = read_column(version, field.type.child, nodes, buffers, body)
+        field.type.build_array(length, validity, child)
       when StructType
         children = field.type.children.collect do |child|
-          read_column(child, nodes, buffers, body)
+          read_column(version, child, nodes, buffers, body)
         end
         field.type.build_array(length, validity, children)
       when DenseUnionType
-        # dense union type doesn't have validity.
-        types = validity
+        if version == FB::MetadataVersion::V4
+          # Dense union type has validity with V4.
+          types_buffer = buffers.shift
+          types = body&.slice(types_buffer.offset, types_buffer.length)
+        else
+          # Dense union type doesn't have validity.
+          types = validity
+        end
         offsets_buffer = buffers.shift
-        offsets = body.slice(offsets_buffer.offset, offsets_buffer.length)
+        offsets = body&.slice(offsets_buffer.offset, offsets_buffer.length)
         children = field.type.children.collect do |child|
-          read_column(child, nodes, buffers, body)
+          read_column(version, child, nodes, buffers, body)
         end
         field.type.build_array(length, types, offsets, children)
       when SparseUnionType
-        # sparse union type doesn't have validity.
-        types = validity
+        if version == FB::MetadataVersion::V4
+          # Sparse union type has validity with V4.
+          types_buffer = buffers.shift
+          types = body&.slice(types_buffer.offset, types_buffer.length)
+        else
+          # Sparse union type doesn't have validity.
+          types = validity
+        end
         children = field.type.children.collect do |child|
-          read_column(child, nodes, buffers, body)
+          read_column(version, child, nodes, buffers, body)
         end
         field.type.build_array(length, types, children)
       when DictionaryType
         indices_buffer = buffers.shift
-        indices = body.slice(indices_buffer.offset, indices_buffer.length)
-        dictionary = find_dictionary(field.dictionary_id)
-        field.type.build_array(length, validity, indices, dictionary)
+        indices = body&.slice(indices_buffer.offset, indices_buffer.length)
+        dictionaries = find_dictionaries(field.type.id)
+        field.type.build_array(length, validity, indices, dictionaries)
       end
     end
   end

data/lib/arrow-format/record-batch.rb

@@ -14,15 +14,27 @@
 # specific language governing permissions and limitations
 # under the License.
 
+require_relative "buffer-alignable"
+
 module ArrowFormat
   class RecordBatch
+    include BufferAlignable
+
     attr_reader :schema
     attr_reader :n_rows
+    alias_method :size, :n_rows
+    alias_method :length, :n_rows
     attr_reader :columns
-    def initialize(schema, n_rows, columns)
+    attr_reader :message_metadata
+    def initialize(schema, n_rows, columns, message_metadata: nil)
       @schema = schema
       @n_rows = n_rows
       @columns = columns
+      @message_metadata = message_metadata
+    end
+
+    def empty?
+      @n_rows.zero?
     end
 
     def to_h
@@ -32,5 +44,62 @@ module ArrowFormat
       end
       hash
     end
+
+    def to_flatbuffers
+      fb_record_batch = FB::RecordBatch::Data.new
+      fb_record_batch.length = @n_rows
+      fb_record_batch.nodes = all_columns_enumerator.collect do |array|
+        field_node = FB::FieldNode::Data.new
+        field_node.length = array.size
+        field_node.null_count = array.n_nulls
+        field_node
+      end
+      offset = 0
+      fb_record_batch.buffers = all_buffers_enumerator.collect do |buffer|
+        fb_buffer = FB::Buffer::Data.new
+        fb_buffer.offset = offset
+        if buffer
+          aligned_size = aligned_buffer_size(buffer)
+          offset += aligned_size
+          fb_buffer.length = aligned_size
+        else
+          fb_buffer.length = 0
+        end
+        fb_buffer
+      end
+      # body_compression = FB::BodyCompression::Data.new
+      # body_compression.codec = ...
+      # fb_record_batch.compression = body_compression
+      fb_record_batch
+    end
+
+    # Pre-order depth-first traversal
+    def all_columns_enumerator
+      Enumerator.new do |yielder|
+        traverse = lambda do |array|
+          yielder << array
+          if array.respond_to?(:child)
+            traverse.call(array.child)
+          elsif array.respond_to?(:children)
+            array.children.each do |child_array|
+              traverse.call(child_array)
+            end
+          end
+        end
+        @columns.each do |array|
+          traverse.call(array)
+        end
+      end
+    end
+
+    def all_buffers_enumerator
+      Enumerator.new do |yielder|
+        all_columns_enumerator.each do |array|
+          array.each_buffer do |buffer|
+            yielder << buffer
+          end
+        end
+      end
+    end
   end
 end

data/lib/arrow-format/schema.rb

@@ -17,8 +17,21 @@
 module ArrowFormat
   class Schema
     attr_reader :fields
-    def initialize(fields)
+    attr_reader :metadata
+    attr_reader :message_metadata
+    def initialize(fields, metadata: nil, message_metadata: nil)
       @fields = fields
+      @metadata = metadata
+      @message_metadata = message_metadata
+    end
+
+    def to_flatbuffers
+      fb_schema = FB::Schema::Data.new
+      fb_schema.endianness = FB::Endianness::LITTLE
+      fb_schema.fields = fields.collect(&:to_flatbuffers)
+      fb_schema.custom_metadata = FB.build_custom_metadata(@metadata)
+      # fb_schema.features = @features
+      fb_schema
     end
   end
 end

data/lib/arrow-format/streaming-pull-reader.rb

@@ -16,6 +16,7 @@
 # under the License.
 
 require_relative "array"
+require_relative "dictionary"
 require_relative "error"
 require_relative "field"
 require_relative "readable"
@@ -100,11 +101,23 @@ module ArrowFormat
     private
     def consume_initial(target)
       continuation = target.get_value(CONTINUATION_TYPE, 0)
-      unless continuation == CONTINUATION_INT32
+      if continuation == CONTINUATION_INT32
+        @state = :metadata_length
+      elsif continuation < 0
         raise ReadError.new("Invalid continuation token: " +
                             continuation.inspect)
+      else
+        # For backward compatibility of data produced prior to version
+        # 0.15.0. It doesn't have continuation token. Ignore it and
+        # re-read it as metadata length.
+        metadata_length = continuation
+        if metadata_length == 0
+          @state = :eos
+        else
+          @metadata_length = metadata_length
+          @state = :metadata
+        end
       end
-      @state = :metadata_length
     end
 
     def consume_metadata_length(target)
@@ -123,7 +136,7 @@ module ArrowFormat
 
     def consume_metadata(target)
       metadata_buffer = target.slice(0, @metadata_length)
-      @message = Org::Apache::Arrow::Flatbuf::Message.new(metadata_buffer)
+      @message = FB::Message.new(metadata_buffer)
       @body_length = @message.body_length
       if @body_length < 0
         raise ReadError.new("Negative body length: " +
@@ -174,7 +187,7 @@ module ArrowFormat
         process_schema_message(message, body)
       when :initial_dictionaries
         header = message.header
-        unless header.is_a?(Org::Apache::Arrow::Flatbuf::DictionaryBatch)
+        unless header.is_a?(FB::DictionaryBatch)
           raise ReadError.new("Not a dictionary batch message: " +
                               header.inspect)
         end
@@ -184,9 +197,9 @@ module ArrowFormat
         end
       when :data
         case message.header
-        when Org::Apache::Arrow::Flatbuf::DictionaryBatch
+        when FB::DictionaryBatch
           process_dictionary_batch_message(message, body)
-        when Org::Apache::Arrow::Flatbuf::RecordBatch
+        when FB::RecordBatch
           process_record_batch_message(message, body)
         end
       end
@@ -194,17 +207,17 @@ module ArrowFormat
 
     def process_schema_message(message, body)
       header = message.header
-      unless header.is_a?(Org::Apache::Arrow::Flatbuf::Schema)
+      unless header.is_a?(FB::Schema)
         raise ReadError.new("Not a schema message: " +
                             header.inspect)
       end
 
-      @schema = read_schema(header)
+      @schema = read_schema(header, message.custom_metadata)
       @dictionaries = {}
       @dictionary_fields = {}
       @schema.fields.each do |field|
         next unless field.type.is_a?(DictionaryType)
-        @dictionary_fields[field.dictionary_id] = field
+        @dictionary_fields[field.type.id] = field
       end
       if @dictionaries.size < @dictionary_fields.size
         @state = :initial_dictionaries
@@ -213,31 +226,44 @@ module ArrowFormat
       end
     end
 
-    def process_dictionary_batch_message(message, body)
-      header = message.header
-      if @state == :initial_dictionaries and header.delta?
+    def process_dictionary_batch_message(fb_message, body)
+      fb_header = fb_message.header
+      if @state == :initial_dictionaries and fb_header.delta?
         raise ReadError.new("An initial dictionary batch message must be " +
                             "a non delta dictionary batch message: " +
-                            header.inspect)
+                            fb_header.inspect)
       end
-      field = @dictionary_fields[header.id]
+      field = @dictionary_fields[fb_header.id]
       value_type = field.type.value_type
-      schema = Schema.new([Field.new("dummy", value_type, true, nil)])
-      record_batch = read_record_batch(header.data, schema, body)
-      if header.delta?
-        @dictionaries[header.id] << record_batch.columns[0]
+      schema = Schema.new([Field.new("dummy", value_type)])
+      record_batch = read_record_batch(fb_message.version,
+                                       fb_header.data,
+                                       nil,
+                                       schema,
+                                       body)
+      message_metadata = read_custom_metadata(fb_message.custom_metadata)
+      dictionary = Dictionary.new(fb_header.id,
+                                  record_batch.columns[0],
+                                  message_metadata: message_metadata)
+      if fb_header.delta?
+        @dictionaries[fb_header.id] << dictionary
       else
-        @dictionaries[header.id] = [record_batch.columns[0]]
+        @dictionaries[fb_header.id] = [dictionary]
       end
     end
 
-    def find_dictionary(id)
+    def find_dictionaries(id)
       @dictionaries[id]
     end
 
-    def process_record_batch_message(message, body)
-      header = message.header
-      @on_read.call(read_record_batch(header, @schema, body))
+    def process_record_batch_message(fb_message, body)
+      fb_header = fb_message.header
+      record_batch = read_record_batch(fb_message.version,
+                                       fb_header,
+                                       fb_message.custom_metadata,
+                                       @schema,
+                                       body)
+      @on_read.call(record_batch)
     end
   end
 end