red-arrow-format 23.0.1 → 24.0.0

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

@@ -21,29 +21,66 @@ module ArrowFormat
   class StreamingReader
     include Enumerable
 
-    attr_reader :schema
     def initialize(input)
-      @input = input
-      @schema = nil
+      case input
+      when File
+        @input = IO::Buffer.map(input, nil, 0, IO::Buffer::READONLY)
+        @offset = 0
+      when String
+        @input = IO::Buffer.for(input)
+        @offset = 0
+      else
+        @input = input
+      end
+
+      @on_read = nil
+      @pull_reader = StreamingPullReader.new do |record_batch|
+        @on_read.call(record_batch) if @on_read
+      end
+      @buffer = "".b
+      ensure_schema
     end
 
-    def each
+    def schema
+      @pull_reader.schema
+    end
+
+    def each(&block)
       return to_enum(__method__) unless block_given?
 
-      reader = StreamingPullReader.new do |record_batch|
-        @schema ||= reader.schema
-        yield(record_batch)
+      @on_read = block
+      begin
+        loop do
+          break unless consume
+        end
+      ensure
+        @on_read = nil
       end
+    end
 
-      buffer = "".b
-      loop do
-        next_size = reader.next_required_size
-        break if next_size.zero?
+    private
+    def consume
+      next_size = @pull_reader.next_required_size
+      return false if next_size.zero?
 
-        next_chunk = @input.read(next_size, buffer)
-        break if next_chunk.nil?
+      if @input.is_a?(IO::Buffer)
+        next_chunk = @input.slice(@offset, next_size)
+        @offset += next_size
+        @pull_reader.consume(next_chunk)
+        true
+      else
+        next_chunk = @input.read(next_size, @buffer)
+        return false if next_chunk.nil?
 
-        reader.consume(IO::Buffer.for(next_chunk))
+        @pull_reader.consume(IO::Buffer.for(next_chunk))
+        true
+      end
+    end
+
+    def ensure_schema
+      loop do
+        break unless consume
+        break if @pull_reader.schema
       end
     end
   end

data/lib/arrow-format/streaming-writer.rb

@@ -0,0 +1,156 @@
+# 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 StreamingWriter
+    include BufferAlignable
+
+    ALIGNMENT_SIZE = IO::Buffer.size_of(:u64)
+    CONTINUATION = "\xFF\xFF\xFF\xFF".b.freeze
+    EOS = "\xFF\xFF\xFF\xFF\x00\x00\x00\x00".b.freeze
+    METADATA_LARGEST_PADDING = "\x00" * 7
+
+    def initialize(output)
+      @output = output
+      @offset = 0
+      @fb_dictionary_blocks = []
+      @fb_record_batch_blocks = []
+      @written_dictionary_offsets = {}
+    end
+
+    def start(schema)
+      write_message(build_metadata(schema.to_flatbuffers,
+                                   custom_metadata: schema.message_metadata))
+    end
+
+    def write_record_batch(record_batch)
+      record_batch.schema.fields.each_with_index do |field, i|
+        next unless field.type.is_a?(DictionaryType)
+        dictionary_array = record_batch.columns[i]
+        write_dictionary(field.type.id, dictionary_array)
+      end
+
+      write_record_batch_based_message(record_batch,
+                                       record_batch.to_flatbuffers,
+                                       @fb_record_batch_blocks)
+    end
+
+    def finish
+      write_data(EOS)
+      @output
+    end
+
+    private
+    def write_data(data)
+      case data
+      when IO::Buffer
+        # TODO: We should use IO::Buffer#write to avoid needless copy.
+        # data.write(@output)
+        @output << data.get_string
+        @offset += data.size
+      else
+        @output << data
+        @offset += data.bytesize
+      end
+    end
+
+    def write_buffer(buffer)
+      write_data(buffer)
+      padding_size = buffer_padding_size(buffer)
+      write_data(padding(padding_size)) if padding_size > 0
+    end
+
+    def build_metadata(header, body_length=0, custom_metadata: nil)
+      fb_message = FB::Message::Data.new
+      fb_message.version = FB::MetadataVersion::V5
+      fb_message.header = header
+      fb_message.body_length = body_length
+      fb_message.custom_metadata = FB.build_custom_metadata(custom_metadata)
+      metadata = FB::Message.serialize(fb_message)
+      metadata_size = metadata.bytesize
+      padding_size = compute_padding_size(metadata_size, ALIGNMENT_SIZE)
+      metadata_size += padding_size
+      align!(metadata, ALIGNMENT_SIZE)
+      metadata
+    end
+
+    def write_record_batch_based_message(record_batch, fb_header, fb_blocks)
+      body_length = 0
+      record_batch.all_buffers_enumerator.each do |buffer|
+        body_length += aligned_buffer_size(buffer) if buffer
+      end
+      metadata = build_metadata(fb_header, body_length,
+                                custom_metadata: record_batch.message_metadata)
+      fb_block = FB::Block::Data.new
+      fb_block.offset = @offset
+      fb_block.meta_data_length =
+        CONTINUATION.bytesize +
+        MessagePullReader::METADATA_LENGTH_SIZE +
+        metadata.bytesize
+      fb_block.body_length = body_length
+      fb_blocks << fb_block
+      write_message(metadata) do
+        record_batch.all_buffers_enumerator.each do |buffer|
+          write_buffer(buffer) if buffer
+        end
+      end
+    end
+
+    def write_dictionary(id, dictionary_array)
+      value_type = dictionary_array.type.value_type
+      base_offset = 0
+      dictionary_array.dictionaries.each do |dictionary|
+        data = dictionary.array
+        written_offset = @written_dictionary_offsets[id] || 0
+        current_base_offset = base_offset
+        next_base_offset = base_offset + data.size
+        base_offset = next_base_offset
+
+        next if next_base_offset <= written_offset
+
+        is_delta = (not written_offset.zero?)
+        if current_base_offset < written_offset
+          data = data.slice(written_offset - current_base_offset)
+        end
+
+        schema = Schema.new([Field.new("dummy", value_type)])
+        size = data.size
+        record_batch =
+          RecordBatch.new(schema, size, [data],
+                          message_metadata: dictionary.message_metadata)
+        fb_dictionary_batch = FB::DictionaryBatch::Data.new
+        fb_dictionary_batch.id = id
+        fb_dictionary_batch.data = record_batch.to_flatbuffers
+        fb_dictionary_batch.delta = is_delta
+        write_record_batch_based_message(record_batch,
+                                         fb_dictionary_batch,
+                                         @fb_dictionary_blocks)
+        @written_dictionary_offsets[id] = written_offset + data.size
+      end
+    end
+
+    def write_message(metadata)
+      write_data(CONTINUATION)
+      metadata_size = metadata.bytesize
+      write_data([metadata_size].pack("l<"))
+      write_data(metadata)
+      yield if block_given?
+    end
+  end
+end