purplelight 0.1.13 → 0.1.17

data/lib/purplelight/writer_parquet.rb

@@ -27,7 +27,7 @@ module Purplelight
       @rotate_rows = rotate_rows
 
       @closed = false
-      @file_seq = 0
+      @file_seq = manifest ? manifest.parts.length : 0
       @part_index = nil
       @pq_writer = nil
       @rows_in_current_file = 0
@@ -37,8 +37,7 @@ module Purplelight
     end
 
     def write_many(array_of_docs)
-      ensure_open!
-      array_of_docs.each { |doc| @buffer_docs << doc }
+      @buffer_docs.concat(array_of_docs)
       flush_row_groups_if_needed
     end
 
@@ -61,6 +60,9 @@ module Purplelight
     def reset_buffers
       @buffer_docs = []
       @columns = nil
+      @buffer_head = 0
+      @arrow_schema = nil
+      @arrow_data_types = nil
       @writer_path = nil
     end
 
@@ -69,57 +71,235 @@ module Purplelight
 
       FileUtils.mkdir_p(@directory)
       @writer_path = next_part_path
-      @part_index = @manifest&.open_part!(@writer_path) if @manifest
+      @part_index = @manifest.open_part!(@writer_path) if @manifest
       @rows_in_current_file = 0
     end
 
     # No-op; we now write once on close for simplicity
 
-    def build_table(docs)
-      # Infer columns
-      @columns ||= infer_columns(docs)
-      columns = {}
-      @columns.each do |name|
-        values = docs.map { |d| extract_value(d, name) }
-        columns[name] = Arrow::ArrayBuilder.build(values)
+    def build_record_batch(documents)
+      @columns ||= infer_columns(documents)
+      arrays = @columns.each_with_index.map do |name, index|
+        values = documents.map { |document| extract_value(document, name) }
+        build_arrow_array(values, @arrow_data_types&.at(index))
+      end
+      batch = if @arrow_schema
+                Arrow::RecordBatch.new(@arrow_schema, arrays)
+              else
+                Arrow::RecordBatch.new(@columns.zip(arrays).to_h)
+              end
+      unless @arrow_schema
+        @arrow_schema = batch.schema
+        @arrow_data_types = @arrow_schema.fields.map(&:data_type)
       end
-      Arrow::Table.new(columns)
+      batch
     end
 
-    def write_table(table, path, append: false) # rubocop:disable Lint/UnusedMethodArgument
-      # Stream via ArrowFileWriter when available to avoid building huge tables
-      if defined?(Parquet::ArrowFileWriter)
-        unless @pq_writer
-          props = build_writer_properties_for_compression(@compression)
-          @pq_writer = create_arrow_file_writer(table.schema, path, props)
-        end
-        # Prefer passing row_group_size; fallback to single-arg for older APIs
-        begin
-          @pq_writer.write_table(table, @row_group_size)
-        rescue ArgumentError
-          @pq_writer.write_table(table)
-        end
-        return
+    def build_arrow_array(values, data_type = nil)
+      first = values.find { |value| !value.nil? }
+      if first.is_a?(BSON::ObjectId)
+        return build_object_id_array(values) if values.all? { |value| value.nil? || value.is_a?(BSON::ObjectId) }
+
+        return build_string_array(values)
       end
-      # Fallback to one-shot save when streaming API is not available
-      if table.respond_to?(:save)
-        begin
-          table.save(path, format: :parquet, compression: normalize_parquet_compression_name(@compression))
-        rescue StandardError
-          table.save(path, format: :parquet)
-        end
-        return
+      return build_string_array(values) if data_type.is_a?(Arrow::StringDataType) ||
+                                           first.is_a?(String) || first.is_a?(Hash)
+      if (data_type.is_a?(Arrow::ListDataType) || first.is_a?(Array)) &&
+         values.all? { |value| value.nil? || value.is_a?(Array) } &&
+         (data_type || values.any? { |value| value && !value.empty? })
+        return build_list_array(values, data_type)
+      end
+      if (data_type.is_a?(Arrow::BooleanDataType) || first.equal?(true) || first.equal?(false)) &&
+         values.all? { |value| value.nil? || value.equal?(true) || value.equal?(false) }
+        return build_boolean_array(values)
+      end
+
+      integer_spec = integer_array_spec(data_type)
+      integer_values = (integer_spec || first.is_a?(Integer)) &&
+                       values.all? { |value| value.nil? || value.is_a?(Integer) }
+      return build_integer_array(values, data_type, integer_spec) if integer_values
+      if (data_type.is_a?(Arrow::DoubleDataType) || first.is_a?(Float)) &&
+         values.all? { |value| value.nil? || value.is_a?(Float) }
+        return build_double_array(values)
       end
-      raise 'Parquet writer not available in this environment'
+
+      data_type ? data_type.build_array(values) : Arrow::ArrayBuilder.build(values)
     end
 
-    def finalize_current_part!
-      return if @writer_path.nil?
+    def build_string_array(values)
+      offsets = Array.new(values.length + 1, 0)
+      data = +''.b
+      values.each_with_index do |value, index|
+        data << (value.is_a?(String) ? value : value.to_s) unless value.nil?
+        offsets[index + 1] = data.bytesize
+      end
+      validity, null_count = build_validity(values)
+      Arrow::StringArray.new(
+        values.length,
+        Arrow::Buffer.new(offsets.pack('l<*').freeze),
+        Arrow::Buffer.new(data.freeze),
+        validity,
+        null_count
+      )
+    end
+
+    def build_object_id_array(values)
+      offsets = Array.new(values.length + 1, 0)
+      raw = String.new(capacity: values.length * 12, encoding: Encoding::BINARY)
+      values.each_with_index do |value, index|
+        # #marshal_dump is BSON's public accessor for an ObjectId's 12 raw bytes.
+        raw << value.marshal_dump if value
+        offsets[index + 1] = raw.bytesize * 2
+      end
+      data = raw.unpack1('H*').force_encoding(Encoding::UTF_8)
+      validity, null_count = build_validity(values)
+      Arrow::StringArray.new(
+        values.length,
+        Arrow::Buffer.new(offsets.pack('l<*').freeze),
+        Arrow::Buffer.new(data.freeze),
+        validity,
+        null_count
+      )
+    end
+
+    def build_list_array(values, data_type)
+      offsets = Array.new(values.length + 1, 0)
+      flattened = []
+      values.each_with_index do |value, index|
+        flattened.concat(value) if value
+        offsets[index + 1] = flattened.length
+      end
+      child_type = data_type&.field&.data_type
+      child = build_arrow_array(flattened, child_type)
+      data_type ||= Arrow::ListDataType.new(child.value_data_type)
+      validity, null_count = build_validity(values)
+      Arrow::ListArray.new(
+        data_type,
+        values.length,
+        Arrow::Buffer.new(offsets.pack('l<*').freeze),
+        child,
+        validity,
+        null_count
+      )
+    end
+
+    def build_boolean_array(values)
+      data = "\0".b * ((values.length + 7) / 8)
+      values.each_with_index do |value, index|
+        next unless value
+
+        byte_index = index >> 3
+        data.setbyte(byte_index, data.getbyte(byte_index) | (1 << (index & 7)))
+      end
+      validity, null_count = build_validity(values)
+      Arrow::BooleanArray.new(values.length, Arrow::Buffer.new(data.freeze), validity, null_count)
+    end
+
+    def build_integer_array(values, data_type, spec)
+      non_null_values = values.compact
+      minimum, maximum = non_null_values.minmax
+      data_type ||= infer_integer_data_type(minimum, maximum)
+      spec ||= integer_array_spec(data_type)
+      return Arrow::ArrayBuilder.build(values) unless spec
+
+      array_class, pack_directive, lower_bound, upper_bound = spec
+      outside_bounds = minimum && (minimum < lower_bound || maximum > upper_bound)
+      raise RangeError, "integer range #{minimum}..#{maximum} exceeds #{data_type}" if outside_bounds
+
+      packed_values = if non_null_values.length == values.length
+                        values
+                      else
+                        values.map { |value| value.nil? ? 0 : value }
+                      end
+      validity, null_count = build_validity(values)
+      array_class.new(
+        values.length,
+        Arrow::Buffer.new(packed_values.pack(pack_directive).freeze),
+        validity,
+        null_count
+      )
+    end
+
+    def build_double_array(values)
+      packed_values = if values.include?(nil)
+                        values.map { |value| value.nil? ? 0.0 : value }
+                      else
+                        values
+                      end
+      validity, null_count = build_validity(values)
+      Arrow::DoubleArray.new(
+        values.length,
+        Arrow::Buffer.new(packed_values.pack('E*').freeze),
+        validity,
+        null_count
+      )
+    end
+
+    def build_validity(values)
+      null_count = values.count(nil)
+      return [nil, 0] if null_count.zero?
+
+      bytes = "\0".b * ((values.length + 7) / 8)
+      values.each_with_index do |value, index|
+        next if value.nil?
+
+        byte_index = index >> 3
+        bytes.setbyte(byte_index, bytes.getbyte(byte_index) | (1 << (index & 7)))
+      end
+      [Arrow::Buffer.new(bytes.freeze), null_count]
+    end
 
-      if @pq_writer
-        @pq_writer.close
-        @pq_writer = nil
+    def infer_integer_data_type(minimum, maximum)
+      if minimum.negative?
+        return Arrow::Int8DataType.new if minimum >= -128 && maximum <= 127
+        return Arrow::Int16DataType.new if minimum >= -32_768 && maximum <= 32_767
+        return Arrow::Int32DataType.new if minimum >= -2_147_483_648 && maximum <= 2_147_483_647
+
+        int64_range = minimum >= -9_223_372_036_854_775_808 && maximum <= 9_223_372_036_854_775_807
+        return Arrow::Int64DataType.new if int64_range
+      else
+        return Arrow::UInt8DataType.new if maximum <= 255
+        return Arrow::UInt16DataType.new if maximum <= 65_535
+        return Arrow::UInt32DataType.new if maximum <= 4_294_967_295
+        return Arrow::UInt64DataType.new if maximum <= 18_446_744_073_709_551_615
       end
+      nil
+    end
+
+    def integer_array_spec(data_type)
+      case data_type
+      when Arrow::UInt8DataType
+        [Arrow::UInt8Array, 'C*', 0, 255]
+      when Arrow::UInt16DataType
+        [Arrow::UInt16Array, 'S<*', 0, 65_535]
+      when Arrow::UInt32DataType
+        [Arrow::UInt32Array, 'L<*', 0, 4_294_967_295]
+      when Arrow::UInt64DataType
+        [Arrow::UInt64Array, 'Q<*', 0, 18_446_744_073_709_551_615]
+      when Arrow::Int8DataType
+        [Arrow::Int8Array, 'c*', -128, 127]
+      when Arrow::Int16DataType
+        [Arrow::Int16Array, 's<*', -32_768, 32_767]
+      when Arrow::Int32DataType
+        [Arrow::Int32Array, 'l<*', -2_147_483_648, 2_147_483_647]
+      when Arrow::Int64DataType
+        [Arrow::Int64Array, 'q<*', -9_223_372_036_854_775_808, 9_223_372_036_854_775_807]
+      end
+    end
+
+    def write_record_batch(batch, path)
+      unless @pq_writer
+        properties = build_writer_properties_for_compression(@compression)
+        properties ||= Parquet::WriterProperties.new
+        properties.max_row_group_length = @row_group_size
+        @pq_writer = Parquet::ArrowFileWriter.open(batch.schema, path, properties)
+      end
+      @pq_writer.write(batch)
+    end
+
+    def finalize_current_part!
+      @pq_writer.close
+      @pq_writer = nil
       @manifest&.complete_part!(index: @part_index, checksum: nil)
       @file_seq += 1 unless @single_file
       @writer_path = nil
@@ -146,77 +326,69 @@ module Purplelight
     end
 
     def extract_value(doc, key)
-      value = doc[key] || doc[key.to_sym]
-      # Normalize common MongoDB/BSON types to Parquet-friendly values
-      return value.to_s if defined?(BSON) && value.is_a?(BSON::ObjectId)
-
+      value = doc[key]
+      value = doc[key.to_sym] if value.nil? && !doc.key?(key)
       value
     end
 
     def flush_row_groups_if_needed
-      return if @buffer_docs.empty?
+      return if buffered_count.zero?
 
-      while @buffer_docs.length >= @row_group_size
+      while buffered_count >= @row_group_size
         ensure_open!
-        group = @buffer_docs.shift(@row_group_size)
+        group = take_buffered(@row_group_size)
         if @rotate_rows && !@single_file && (@rows_in_current_file + group.length) > @rotate_rows
           # Write a partial chunk to fill the current file, then rotate and write the rest
           remaining_allowed = @rotate_rows - @rows_in_current_file
-          if remaining_allowed.positive?
-            part_a = group.first(remaining_allowed)
-            t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
-            table_a = build_table(part_a)
-            Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
-
-            t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
-            write_table(table_a, @writer_path, append: true)
-            Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
-            @manifest&.add_progress_to_part!(index: @part_index, rows_delta: part_a.length, bytes_delta: 0)
-            @rows_in_current_file += part_a.length
-          end
+          part_a = group.first(remaining_allowed)
+          t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
+          batch_a = build_record_batch(part_a)
+          Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
+
+          t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
+          write_record_batch(batch_a, @writer_path)
+          Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
+          @manifest&.add_progress_to_part!(index: @part_index, rows_delta: part_a.length, bytes_delta: 0)
+          @rows_in_current_file += part_a.length
 
           finalize_current_part!
           ensure_open!
 
           part_b = group.drop(remaining_allowed)
-          unless part_b.empty?
-            t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
-            table_b = build_table(part_b)
-            Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
-
-            t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
-            write_table(table_b, @writer_path, append: true)
-            Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
-            @manifest&.add_progress_to_part!(index: @part_index, rows_delta: part_b.length, bytes_delta: 0)
-            @rows_in_current_file += part_b.length
-            maybe_rotate!
-          end
+          t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
+          batch_b = build_record_batch(part_b)
+          Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
+
+          t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
+          write_record_batch(batch_b, @writer_path)
+          Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
+          @manifest&.add_progress_to_part!(index: @part_index, rows_delta: part_b.length, bytes_delta: 0)
+          @rows_in_current_file += part_b.length
         else
           t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
-          table = build_table(group)
+          batch = build_record_batch(group)
           Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
 
           t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
-          write_table(table, @writer_path, append: true)
+          write_record_batch(batch, @writer_path)
           Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
           @manifest&.add_progress_to_part!(index: @part_index, rows_delta: group.length, bytes_delta: 0)
           @rows_in_current_file += group.length
-          maybe_rotate!
         end
+        maybe_rotate!
       end
     end
 
     def flush_all_row_groups
-      return if @buffer_docs.empty?
+      return if buffered_count.zero?
 
       # Flush any full groups first
       flush_row_groups_if_needed
-      return if @buffer_docs.empty?
 
       # Flush remaining as a final smaller group
-      remaining = @buffer_docs.length
+      remaining = buffered_count
       t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
-      table = build_table(@buffer_docs)
+      batch = build_record_batch(take_buffered(remaining))
       Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
 
       ensure_open!
@@ -227,12 +399,13 @@ module Purplelight
       end
 
       t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
-      write_table(table, @writer_path, append: true)
+      write_record_batch(batch, @writer_path)
       Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
-      rows_written = (table.respond_to?(:n_rows) ? table.n_rows : remaining)
+      rows_written = batch.n_rows
       @manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows_written, bytes_delta: 0)
       @rows_in_current_file += rows_written
       @buffer_docs.clear
+      @buffer_head = 0
       maybe_rotate!
     end
 
@@ -245,76 +418,26 @@ module Purplelight
     end
 
     def build_writer_properties_for_compression(requested)
-      codec_const = parquet_codec_constant(requested)
-      return nil unless codec_const
-
-      # Prefer WriterProperties builder if available
-      begin
-        if defined?(Parquet::WriterProperties) && Parquet::WriterProperties.respond_to?(:builder)
-          builder = Parquet::WriterProperties.builder
-          if builder.respond_to?(:compression)
-            builder = builder.compression(codec_const)
-          elsif builder.respond_to?(:set_compression)
-            builder = builder.set_compression(codec_const)
-          end
-          return builder.build if builder.respond_to?(:build)
-        end
-      rescue StandardError
-        # fall through to other strategies
-      end
+      compression = normalize_parquet_compression_name(requested)
+      return nil unless compression
 
-      # Alternative builder class naming fallback
-      begin
-        if defined?(Parquet::WriterPropertiesBuilder)
-          b = Parquet::WriterPropertiesBuilder.new
-          if b.respond_to?(:compression)
-            b.compression(codec_const)
-          elsif b.respond_to?(:set_compression)
-            b.set_compression(codec_const)
-          end
-          return b.build if b.respond_to?(:build)
-        end
-      rescue StandardError
-        # ignore
-      end
-      nil
+      properties = Parquet::WriterProperties.new
+      properties.set_compression(compression == 'none' ? 'uncompressed' : compression)
+      properties
     end
 
-    def create_arrow_file_writer(schema, path, props)
-      attempts = []
-      if props
-        attempts << -> { Parquet::ArrowFileWriter.open(schema, path, props) }
-        attempts << -> { Parquet::ArrowFileWriter.open(schema, path, properties: props) }
-      end
-      attempts << -> { Parquet::ArrowFileWriter.open(schema, path) }
-
-      attempts.each do |call|
-        return call.call
-      rescue StandardError
-        next
-      end
-      raise 'failed to open Parquet::ArrowFileWriter'
+    def buffered_count
+      @buffer_docs.length - @buffer_head
     end
 
-    def parquet_codec_constant(requested)
-      name = normalize_parquet_compression_name(requested)
-      return nil unless name
-
-      up = case name
-           when 'zstd', 'zstandard' then 'ZSTD'
-           when 'gzip' then 'GZIP'
-           when 'snappy' then 'SNAPPY'
-           when 'none' then 'UNCOMPRESSED'
-           else name.upcase
-           end
-      candidates = %w[CompressionType Compression CompressionCodec]
-      candidates.each do |mod|
-        m = Parquet.const_get(mod)
-        return m.const_get(up) if m.const_defined?(up)
-      rescue StandardError
-        next
+    def take_buffered(count)
+      documents = @buffer_docs.slice(@buffer_head, count)
+      @buffer_head += count
+      if @buffer_head >= @row_group_size && @buffer_head * 2 >= @buffer_docs.length
+        @buffer_docs = @buffer_docs.drop(@buffer_head)
+        @buffer_head = 0
       end
-      nil
+      documents
     end
 
     def normalize_parquet_compression_name(requested)
@@ -324,7 +447,7 @@ module Purplelight
       return 'none' if s == 'none'
       return 'gzip' if s == 'gzip'
       return 'snappy' if s == 'snappy'
-      return 'zstd' if %w[zstd zstandard].include?(s)
+      return 'zstd' if s == 'zstd'
 
       nil
     end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: purplelight
 version: !ruby/object:Gem::Version
-  version: 0.1.13
+  version: 0.1.17
 platform: ruby
 authors:
 - Alexander Nicholson
@@ -10,25 +10,19 @@ cert_chain: []
 date: 1980-01-02 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
-  name: csv
+  name: bigdecimal
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '3.3'
     - - ">="
       - !ruby/object:Gem::Version
-        version: 3.3.5
+        version: '3.1'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '3.3'
     - - ">="
       - !ruby/object:Gem::Version
-        version: 3.3.5
+        version: '3.1'
 - !ruby/object:Gem::Dependency
   name: logger
   requirement: !ruby/object:Gem::Requirement
@@ -55,20 +49,20 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.21'
+        version: '2.24'
     - - ">="
       - !ruby/object:Gem::Version
-        version: 2.21.3
+        version: 2.24.1
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '2.21'
+        version: '2.24'
     - - ">="
       - !ruby/object:Gem::Version
-        version: 2.21.3
+        version: 2.24.1
 description: High-throughput, resumable snapshots of MongoDB collections with partitioning,
   multi-threaded readers, and size-based sharded outputs.
 email:
@@ -79,8 +73,13 @@ extensions: []
 extra_rdoc_files: []
 files:
 - LICENSE
+- PERFORMANCE_NOTES.md
 - README.md
 - Rakefile
+- benchmark/baseline.json
+- benchmark/harness.rb
+- benchmark/microbench.rb
+- benchmark/profile.rb
 - bin/purplelight
 - lib/purplelight.rb
 - lib/purplelight/errors.rb
@@ -114,7 +113,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.6.7
+rubygems_version: 4.0.6
 specification_version: 4
 summary: Snapshot MongoDB collections efficiently to JSONL/CSV/Parquet
 test_files: []