red-arrow-format 23.0.0

data/lib/arrow-format/file-reader.rb

@@ -0,0 +1,213 @@
+# 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 "streaming-reader"
+
+require_relative "org/apache/arrow/flatbuf/block"
+require_relative "org/apache/arrow/flatbuf/footer"
+
+module ArrowFormat
+  class FileReader
+    include Enumerable
+    include Readable
+
+    MAGIC = "ARROW1".b.freeze
+    MAGIC_BUFFER = IO::Buffer.for(MAGIC)
+    START_MARKER_SIZE = MAGIC_BUFFER.size
+    END_MARKER_SIZE = MAGIC_BUFFER.size
+    # <magic number "ARROW1">
+    # <empty padding bytes [to 8 byte boundary]>
+    STREAMING_FORMAT_START_OFFSET = 8
+    CONTINUATION_BUFFER =
+      IO::Buffer.for(MessagePullReader::CONTINUATION_STRING)
+    FOOTER_SIZE_FORMAT = :s32
+    FOOTER_SIZE_SIZE = IO::Buffer.size_of(FOOTER_SIZE_FORMAT)
+
+    def initialize(input)
+      case input
+      when IO
+        @buffer = IO::Buffer.map(input, nil, 0, IO::Buffer::READONLY)
+      when String
+        @buffer = IO::Buffer.for(input)
+      else
+        @buffer = input
+      end
+
+      validate
+      @footer = read_footer
+      @record_batch_blocks = @footer.record_batches
+      @schema = read_schema(@footer.schema)
+      @dictionaries = read_dictionaries
+    end
+
+    def n_record_batches
+      @record_batch_blocks.size
+    end
+
+    def read(i)
+      fb_message, body = read_block(@record_batch_blocks[i])
+      fb_header = fb_message.header
+      unless fb_header.is_a?(Org::Apache::Arrow::Flatbuf::RecordBatch)
+        raise FileReadError.new(@buffer,
+                                "Not a record batch message: #{i}: " +
+                                fb_header.class.name)
+      end
+      read_record_batch(fb_header, @schema, body)
+    end
+
+    def each
+      return to_enum(__method__) {n_record_batches} unless block_given?
+
+      @record_batch_blocks.size.times do |i|
+        yield(read(i))
+      end
+    end
+
+    private
+    def validate
+      minimum_size = STREAMING_FORMAT_START_OFFSET +
+                     FOOTER_SIZE_SIZE +
+                     END_MARKER_SIZE
+      if @buffer.size < minimum_size
+        raise FileReadError.new(@buffer,
+                                "Input must be larger than or equal to " +
+                                "#{minimum_size}: #{@buffer.size}")
+      end
+
+      start_marker = @buffer.slice(0, START_MARKER_SIZE)
+      if start_marker != MAGIC_BUFFER
+        raise FileReadError.new(@buffer, "No start marker")
+      end
+      end_marker = @buffer.slice(@buffer.size - END_MARKER_SIZE,
+                                 END_MARKER_SIZE)
+      if end_marker != MAGIC_BUFFER
+        raise FileReadError.new(@buffer, "No end marker")
+      end
+    end
+
+    def read_footer
+      footer_size_offset = @buffer.size - END_MARKER_SIZE - FOOTER_SIZE_SIZE
+      footer_size = @buffer.get_value(FOOTER_SIZE_FORMAT, footer_size_offset)
+      footer_data = @buffer.slice(footer_size_offset - footer_size,
+                                  footer_size)
+      Org::Apache::Arrow::Flatbuf::Footer.new(footer_data)
+    end
+
+    def read_block(block)
+      offset = block.offset
+
+      # If we can report property error information, we can use
+      # MessagePullReader here.
+      #
+      # message_pull_reader = MessagePullReader.new do |message, body|
+      #   return read_record_batch(message.header, @schema, body)
+      # end
+      # chunk = @buffer.slice(offset,
+      #                       MessagePullReader::CONTINUATION_SIZE +
+      #                       MessagePullReader::METADATA_LENGTH_SIZE +
+      #                       block.meta_data_length +
+      #                       block.body_length)
+      # message_pull_reader.consume(chunk)
+
+      continuation_size = CONTINUATION_BUFFER.size
+      continuation = @buffer.slice(offset, continuation_size)
+      unless continuation == CONTINUATION_BUFFER
+        raise FileReadError.new(@buffer,
+                                "Invalid continuation: #{i}: " +
+                                continuation.inspect)
+      end
+      offset += continuation_size
+
+      metadata_length_type = MessagePullReader::METADATA_LENGTH_TYPE
+      metadata_length_size = MessagePullReader::METADATA_LENGTH_SIZE
+      metadata_length = @buffer.get_value(metadata_length_type, offset)
+      expected_metadata_length =
+        block.meta_data_length -
+        continuation_size -
+        metadata_length_size
+      unless metadata_length == expected_metadata_length
+        raise FileReadError.new(@buffer,
+                                "Invalid metadata length #{i}: " +
+                                "expected:#{expected_metadata_length} " +
+                                "actual:#{metadata_length}")
+      end
+      offset += metadata_length_size
+
+      metadata = @buffer.slice(offset, metadata_length)
+      fb_message = Org::Apache::Arrow::Flatbuf::Message.new(metadata)
+      offset += metadata_length
+
+      body = @buffer.slice(offset, block.body_length)
+
+      [fb_message, body]
+    end
+
+    def read_dictionaries
+      dictionary_blocks = @footer.dictionaries
+      return nil if dictionary_blocks.nil?
+
+      dictionary_fields = {}
+      @schema.fields.each do |field|
+        next unless field.type.is_a?(DictionaryType)
+        dictionary_fields[field.dictionary_id] = field
+      end
+
+      dictionaries = {}
+      dictionary_blocks.each do |block|
+        fb_message, body = read_block(block)
+        fb_header = fb_message.header
+        unless fb_header.is_a?(Org::Apache::Arrow::Flatbuf::DictionaryBatch)
+          raise FileReadError.new(@buffer,
+                                  "Not a dictionary batch message: " +
+                                  fb_header.inspect)
+        end
+
+        id = fb_header.id
+        if fb_header.delta?
+          unless dictionaries.key?(id)
+            raise FileReadError.new(@buffer,
+                                    "A delta dictionary batch message " +
+                                    "must exist after a non delta " +
+                                    "dictionary batch message: " +
+                                    fb_header.inspect)
+          end
+        else
+          if dictionaries.key?(id)
+            raise FileReadError.new(@buffer,
+                                    "Multiple non delta dictionary batch " +
+                                    "messages for the same ID is invalid: " +
+                                    fb_header.inspect)
+          end
+        end
+
+        value_type = dictionary_fields[id].type.value_type
+        schema = Schema.new([Field.new("dummy", value_type, true, nil)])
+        record_batch = read_record_batch(fb_header.data, schema, body)
+        if fb_header.delta?
+          dictionaries[id] << record_batch.columns[0]
+        else
+          dictionaries[id] = [record_batch.columns[0]]
+        end
+      end
+      dictionaries
+    end
+
+    def find_dictionary(id)
+      @dictionaries[id]
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/binary.rb

@@ -0,0 +1,21 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # Opaque binary data
+          class Binary < ::FlatBuffers::Table
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/binary_view.rb

@@ -0,0 +1,27 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # Logically the same as Binary, but the internal representation uses a view
+          # struct that contains the string length and either the string's entire data
+          # inline (for small strings) or an inlined prefix, an index of another buffer,
+          # and an offset pointing to a slice in that buffer (for non-small strings).
+          #
+          # Since it uses a variable number of data buffers, each Field with this type
+          # must have a corresponding entry in `variadicBufferCounts`.
+          class BinaryView < ::FlatBuffers::Table
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/block.rb

@@ -0,0 +1,38 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //File.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Footer (//File.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          class Block < ::FlatBuffers::Struct
+            # Length of the data (this is aligned so there can be a gap between this and
+            # the metadata).
+            def body_length
+              field_offset = 16
+              @view.unpack_long(field_offset)
+            end
+
+            # Length of the metadata
+            def meta_data_length
+              field_offset = 8
+              @view.unpack_int(field_offset)
+            end
+
+            # Index to the start of the RecordBlock (note this is past the Message header)
+            def offset
+              field_offset = 0
+              @view.unpack_long(field_offset)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/body_compression.rb

@@ -0,0 +1,47 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Message.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+require_relative "../../../apache/arrow/flatbuf/compression_type"
+require_relative "../../../apache/arrow/flatbuf/body_compression_method"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # Optional compression for the memory buffers constituting IPC message
+          # bodies. Intended for use with RecordBatch but could be used for other
+          # message types
+          class BodyCompression < ::FlatBuffers::Table
+            # Compressor library.
+            # For LZ4_FRAME, each compressed buffer must consist of a single frame.
+            def codec
+              field_offset = @view.unpack_virtual_offset(4)
+              if field_offset.zero?
+                enum_value = 0
+              else
+                enum_value = @view.unpack_byte(field_offset)
+              end
+              ::ArrowFormat::Org::Apache::Arrow::Flatbuf::CompressionType.try_convert(enum_value) || enum_value
+            end
+
+            # Indicates the way the record batch body was compressed
+            def method
+              field_offset = @view.unpack_virtual_offset(6)
+              if field_offset.zero?
+                enum_value = 0
+              else
+                enum_value = @view.unpack_byte(field_offset)
+              end
+              ::ArrowFormat::Org::Apache::Arrow::Flatbuf::BodyCompressionMethod.try_convert(enum_value) || enum_value
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/body_compression_method.rb

@@ -0,0 +1,31 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Message.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # Provided for forward compatibility in case we need to support different
+          # strategies for compressing the IPC message body (like whole-body
+          # compression rather than buffer-level) in the future
+          class BodyCompressionMethod < ::FlatBuffers::Enum
+            # Each constituent buffer is first compressed with the indicated
+            # compressor, and then written with the uncompressed length in the first 8
+            # bytes as a 64-bit little-endian signed integer followed by the compressed
+            # buffer bytes (and then padding as required by the protocol). The
+            # uncompressed length may be set to -1 to indicate that the data that
+            # follows is not compressed, which can be useful for cases where
+            # compression does not yield appreciable savings.
+            BUFFER = register("BUFFER", 0)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/bool.rb

@@ -0,0 +1,20 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          class Bool < ::FlatBuffers::Table
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/buffer.rb

@@ -0,0 +1,38 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # ----------------------------------------------------------------------
+          # A Buffer represents a single contiguous memory segment
+          class Buffer < ::FlatBuffers::Struct
+            # The absolute length (in bytes) of the memory buffer. The memory is found
+            # from offset (inclusive) to offset + length (non-inclusive). When building
+            # messages using the encapsulated IPC message, padding bytes may be written
+            # after a buffer, but such padding bytes do not need to be accounted for in
+            # the size here.
+            def length
+              field_offset = 8
+              @view.unpack_long(field_offset)
+            end
+
+            # The relative offset into the shared memory page where the bytes for this
+            # buffer starts
+            def offset
+              field_offset = 0
+              @view.unpack_long(field_offset)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/compression_type.rb

@@ -0,0 +1,22 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Message.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          class CompressionType < ::FlatBuffers::Enum
+            LZ4_FRAME = register("LZ4_FRAME", 0)
+            ZSTD = register("ZSTD", 1)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/date.rb

@@ -0,0 +1,36 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+require_relative "../../../apache/arrow/flatbuf/date_unit"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # Date is either a 32-bit or 64-bit signed integer type representing an
+          # elapsed time since UNIX epoch (1970-01-01), stored in either of two units:
+          #
+          # * Milliseconds (64 bits) indicating UNIX time elapsed since the epoch (no
+          #   leap seconds), where the values are evenly divisible by 86400000
+          # * Days (32 bits) since the UNIX epoch
+          class Date < ::FlatBuffers::Table
+            def unit
+              field_offset = @view.unpack_virtual_offset(4)
+              if field_offset.zero?
+                enum_value = 1
+              else
+                enum_value = @view.unpack_short(field_offset)
+              end
+              ::ArrowFormat::Org::Apache::Arrow::Flatbuf::DateUnit.try_convert(enum_value) || enum_value
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/date_unit.rb

@@ -0,0 +1,22 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          class DateUnit < ::FlatBuffers::Enum
+            DAY = register("DAY", 0)
+            MILLISECOND = register("MILLISECOND", 1)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/decimal.rb

@@ -0,0 +1,48 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # Exact decimal value represented as an integer value in two's
+          # complement. Currently 32-bit (4-byte), 64-bit (8-byte), 
+          # 128-bit (16-byte) and 256-bit (32-byte) integers are used.
+          # The representation uses the endianness indicated in the Schema.
+          class Decimal < ::FlatBuffers::Table
+            # Number of bits per value. The accepted widths are 32, 64, 128 and 256.
+            # We use bitWidth for consistency with Int::bitWidth.
+            def bit_width
+              field_offset = @view.unpack_virtual_offset(8)
+              return 128 if field_offset.zero?
+
+              @view.unpack_int(field_offset)
+            end
+
+            # Total number of decimal digits
+            def precision
+              field_offset = @view.unpack_virtual_offset(4)
+              return 0 if field_offset.zero?
+
+              @view.unpack_int(field_offset)
+            end
+
+            # Number of digits after the decimal point "."
+            def scale
+              field_offset = @view.unpack_virtual_offset(6)
+              return 0 if field_offset.zero?
+
+              @view.unpack_int(field_offset)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/dictionary_batch.rb

@@ -0,0 +1,50 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Message.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+require_relative "../../../apache/arrow/flatbuf/record_batch"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          # For sending dictionary encoding information. Any Field can be
+          # dictionary-encoded, but in this case none of its children may be
+          # dictionary-encoded.
+          # There is one vector / column per dictionary, but that vector / column
+          # may be spread across multiple dictionary batches by using the isDelta
+          # flag
+          class DictionaryBatch < ::FlatBuffers::Table
+            def data
+              field_offset = @view.unpack_virtual_offset(6)
+              return nil if field_offset.zero?
+
+              @view.unpack_table(::ArrowFormat::Org::Apache::Arrow::Flatbuf::RecordBatch, field_offset)
+            end
+
+            def id
+              field_offset = @view.unpack_virtual_offset(4)
+              return 0 if field_offset.zero?
+
+              @view.unpack_long(field_offset)
+            end
+
+            # If isDelta is true the values in the dictionary are to be appended to a
+            # dictionary with the indicated id. If isDelta is false this dictionary
+            # should replace the existing dictionary.
+            def delta?
+              field_offset = @view.unpack_virtual_offset(8)
+              return false if field_offset.zero?
+
+              @view.unpack_bool(field_offset)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/arrow-format/org/apache/arrow/flatbuf/dictionary_encoding.rb

@@ -0,0 +1,64 @@
+# Automatically generated. Don't modify manually.
+#
+# Red FlatBuffers version: 0.0.3
+# Declared by:             //Schema.fbs
+# Rooting type:            org.apache.arrow.flatbuf.Message (//Message.fbs)
+
+require "flatbuffers"
+require_relative "../../../apache/arrow/flatbuf/dictionary_kind"
+require_relative "../../../apache/arrow/flatbuf/int"
+
+module ArrowFormat
+  module Org
+    module Apache
+      module Arrow
+        module Flatbuf
+          class DictionaryEncoding < ::FlatBuffers::Table
+            def dictionary_kind
+              field_offset = @view.unpack_virtual_offset(10)
+              if field_offset.zero?
+                enum_value = 0
+              else
+                enum_value = @view.unpack_short(field_offset)
+              end
+              ::ArrowFormat::Org::Apache::Arrow::Flatbuf::DictionaryKind.try_convert(enum_value) || enum_value
+            end
+
+            # The known dictionary id in the application where this data is used. In
+            # the file or streaming formats, the dictionary ids are found in the
+            # DictionaryBatch messages
+            def id
+              field_offset = @view.unpack_virtual_offset(4)
+              return 0 if field_offset.zero?
+
+              @view.unpack_long(field_offset)
+            end
+
+            # The dictionary indices are constrained to be non-negative integers. If
+            # this field is null, the indices must be signed int32. To maximize
+            # cross-language compatibility and performance, implementations are
+            # recommended to prefer signed integer types over unsigned integer types
+            # and to avoid uint64 indices unless they are required by an application.
+            def index_type
+              field_offset = @view.unpack_virtual_offset(6)
+              return nil if field_offset.zero?
+
+              @view.unpack_table(::ArrowFormat::Org::Apache::Arrow::Flatbuf::Int, field_offset)
+            end
+
+            # By default, dictionaries are not ordered, or the order does not have
+            # semantic meaning. In some statistical, applications, dictionary-encoding
+            # is used to represent ordered categorical data, and we provide a way to
+            # preserve that metadata here
+            def ordered?
+              field_offset = @view.unpack_virtual_offset(8)
+              return false if field_offset.zero?
+
+              @view.unpack_bool(field_offset)
+            end
+          end
+        end
+      end
+    end
+  end
+end