purplelight 0.1.11 → 0.1.13

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 80f6e48231b485750fc65529ada74a07758befba9b324b58ef372f077305a144
-  data.tar.gz: 4d7eed034f90155d2686da45a76caa73928cbe2d080d2031539830f6f4399cfb
+  metadata.gz: c500768ed34e247a92be3979f56c1036314b51003cb1ff7953562f1ca8278677
+  data.tar.gz: f12b305e35b201192e219aec23a0c66c13a7b9eb262c038dbac9d1e2a3c77bc5
 SHA512:
-  metadata.gz: 63f49c8dab688ec5cd922304b472ee448aaada1a3e2c113b7a4ddbe2092f3a3d2a83e1fe396066e87971514d8069831486c1a5ad972e604807c6c3289efd8e31
-  data.tar.gz: 6a281a23a0abf3244045b3e99af606881f3f85bab7eeabd7b4ca94e36c823d2405d7e85b60048ec62e5a7ace63fceddd268e185a831398cc7dbd00c213198961
+  metadata.gz: 0f6959b02fb695a4c28802ec555d7b6ce6e7020d348024d2df8381efbb022632916d2915c688131cf79aa042f05aa3e2ec8780ad866fdb774b45b51f53093d6b
+  data.tar.gz: aa2885e120f57bd492c1a98aee91f52b7c83bee15492ce34a0b6d3dae985e0355ea77489e69dfde2f27fb5b51974263b7f7358a51103697d8222a03cd98a4185

data/README.md

@@ -9,7 +9,7 @@ Purplelight is published on RubyGems: [purplelight on RubyGems](https://rubygems
 Add to your Gemfile:
 
 ```ruby
-gem 'purplelight', '~> 0.1.11'
+gem 'purplelight', '~> 0.1.13'
 ```
 
 Or install directly:
@@ -212,6 +212,7 @@ bundle exec bin/purplelight \
 - `--read-concern LEVEL`: `majority|local|linearizable|available|snapshot`.
 - `--no-cursor-timeout BOOL`: Toggle `noCursorTimeout` (default true).
 - `--parquet-row-group N`: Parquet row group size (rows).
+- `--parquet-max-rows N`: Parquet max rows per part (enables multi-part when not `--single-file`).
 - `--write-chunk-mb MB`: JSONL encode/write chunk size before enqueueing.
 - `--writer-threads N` (experimental): Number of writer threads (JSONL only).
 - `--telemetry on|off`: Force enable/disable telemetry output.
@@ -222,7 +223,7 @@ bundle exec bin/purplelight \
 Notes:
 - Compression backend selection order is: requested format → `zstd-ruby` → `zstds` → `gzip`.
 - `--single-file` and `--by-size` update only the sharding mode/params and preserve any provided `--prefix`.
-- Parquet multi-part sizing is programmatic via `parquet_max_rows`; there is no CLI flag for it.
+– Parquet multi-part sizing can be controlled via `parquet_max_rows` programmatically, or via CLI `--parquet-max-rows`.
 - To increase concurrent connections, set `maxPoolSize` on your Mongo URI (used by `--uri`), e.g., `mongodb://.../?maxPoolSize=32`. A good starting point is `maxPoolSize >= --partitions`.
 
 ### Architecture

data/bin/purplelight

@@ -25,6 +25,7 @@ options = {
   writer_threads: nil,
   write_chunk_bytes: nil,
   parquet_row_group: nil,
+  parquet_max_rows: nil,
   telemetry_flag: nil,
   read_concern: nil,
   no_cursor_timeout: nil,
@@ -61,6 +62,7 @@ parser = OptionParser.new do |opts|
   opts.on('--writer-threads N', Integer, 'Number of writer threads (experimental, JSONL only)') { |v| options[:writer_threads] = v }
   opts.on('--write-chunk-mb MB', Integer, 'JSONL encode/write chunk size in MB') { |v| options[:write_chunk_bytes] = v * 1024 * 1024 }
   opts.on('--parquet-row-group N', Integer, 'Parquet row group size (rows)') { |v| options[:parquet_row_group] = v }
+  opts.on('--parquet-max-rows N', Integer, 'Parquet max rows per part (multi-part only)') { |v| options[:parquet_max_rows] = v }
   opts.on('-q', '--query JSON', 'Filter query as JSON (Extended JSON supported)') do |v|
     # Prefer BSON Extended JSON to support $date, $oid, etc.
     options[:query] = BSON::ExtJSON.parse(v)
@@ -158,7 +160,8 @@ snapshot_args[:no_cursor_timeout] = options[:no_cursor_timeout] unless options[:
 snapshot_args[:compression_level] = options[:compression_level] if options[:compression_level]
 snapshot_args[:writer_threads] = options[:writer_threads] if options[:writer_threads]
 snapshot_args[:write_chunk_bytes] = options[:write_chunk_bytes] if options[:write_chunk_bytes]
-snapshot_args[:parquet_row_group] = options[:parquet_row_group] if options[:parquet_row_group]
+ snapshot_args[:parquet_row_group] = options[:parquet_row_group] if options[:parquet_row_group]
+ snapshot_args[:parquet_max_rows] = options[:parquet_max_rows] if options[:parquet_max_rows]
 
 # telemetry env override
 if options[:telemetry_flag]
@@ -169,6 +172,7 @@ end
 ENV['PL_ZSTD_LEVEL'] = options[:compression_level].to_s if options[:compression_level]
 ENV['PL_WRITE_CHUNK_BYTES'] = options[:write_chunk_bytes].to_s if options[:write_chunk_bytes]
 ENV['PL_PARQUET_ROW_GROUP'] = options[:parquet_row_group].to_s if options[:parquet_row_group]
+  # No env default for parquet_max_rows; it is snapshot-argument only
 ENV['PL_WRITER_THREADS'] = options[:writer_threads].to_s if options[:writer_threads]
 
 ok = Purplelight.snapshot(**snapshot_args)

data/lib/purplelight/snapshot.rb

@@ -137,8 +137,8 @@ module Purplelight
                  single_file = @sharding && @sharding[:mode].to_s == 'single_file'
                  row_group = @parquet_row_group || ENV['PL_PARQUET_ROW_GROUP']&.to_i || WriterParquet::DEFAULT_ROW_GROUP_SIZE
                  WriterParquet.new(directory: dir, prefix: prefix, compression: @compression, logger: @logger,
-                          manifest: manifest, single_file: single_file, row_group_size: row_group,
-                          rotate_rows: @parquet_max_rows)
+                                   manifest: manifest, single_file: single_file, row_group_size: row_group,
+                                   rotate_rows: @parquet_max_rows)
                else
                  raise ArgumentError, "format not implemented: #{@format}"
                end

data/lib/purplelight/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Purplelight
-  VERSION = '0.1.11'
+  VERSION = '0.1.13'
 end

data/lib/purplelight/writer_parquet.rb

@@ -40,7 +40,6 @@ module Purplelight
       ensure_open!
       array_of_docs.each { |doc| @buffer_docs << doc }
       flush_row_groups_if_needed
-      @manifest&.add_progress_to_part!(index: @part_index, rows_delta: array_of_docs.length, bytes_delta: 0)
     end
 
     def close
@@ -91,7 +90,8 @@ module Purplelight
       # Stream via ArrowFileWriter when available to avoid building huge tables
       if defined?(Parquet::ArrowFileWriter)
         unless @pq_writer
-          @pq_writer = Parquet::ArrowFileWriter.open(table.schema, path)
+          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
@@ -103,7 +103,11 @@ module Purplelight
       end
       # Fallback to one-shot save when streaming API is not available
       if table.respond_to?(:save)
-        table.save(path, format: :parquet)
+        begin
+          table.save(path, format: :parquet, compression: normalize_parquet_compression_name(@compression))
+        rescue StandardError
+          table.save(path, format: :parquet)
+        end
         return
       end
       raise 'Parquet writer not available in this environment'
@@ -111,6 +115,7 @@ module Purplelight
 
     def finalize_current_part!
       return if @writer_path.nil?
+
       if @pq_writer
         @pq_writer.close
         @pq_writer = nil
@@ -154,15 +159,50 @@ module Purplelight
       while @buffer_docs.length >= @row_group_size
         ensure_open!
         group = @buffer_docs.shift(@row_group_size)
-        t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
-        table = build_table(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)
-        Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
-        @rows_in_current_file += group.length
-        maybe_rotate!
+        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
+
+          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
+        else
+          t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
+          table = build_table(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)
+          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
       end
     end
 
@@ -174,17 +214,25 @@ module Purplelight
       return if @buffer_docs.empty?
 
       # Flush remaining as a final smaller group
+      remaining = @buffer_docs.length
       t_tbl = Thread.current[:pl_telemetry]&.start(:parquet_table_build_time)
       table = build_table(@buffer_docs)
       Thread.current[:pl_telemetry]&.finish(:parquet_table_build_time, t_tbl)
 
-      t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
       ensure_open!
+      # Pre-rotate to avoid exceeding rotate_rows on this final write
+      if @rotate_rows && !@single_file && @rows_in_current_file.positive? && (@rows_in_current_file + remaining) > @rotate_rows
+        finalize_current_part!
+        ensure_open!
+      end
+
+      t_w = Thread.current[:pl_telemetry]&.start(:parquet_write_time)
       write_table(table, @writer_path, append: true)
       Thread.current[:pl_telemetry]&.finish(:parquet_write_time, t_w)
+      rows_written = (table.respond_to?(:n_rows) ? table.n_rows : remaining)
+      @manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows_written, bytes_delta: 0)
+      @rows_in_current_file += rows_written
       @buffer_docs.clear
-      @rows_in_current_file += table.n_rows if table.respond_to?(:n_rows)
-      @rows_in_current_file += @buffer_docs.length unless table.respond_to?(:n_rows)
       maybe_rotate!
     end
 
@@ -195,5 +243,90 @@ module Purplelight
       finalize_current_part!
       # Next write will open a new part
     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
+
+      # 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
+    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'
+    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
+      end
+      nil
+    end
+
+    def normalize_parquet_compression_name(requested)
+      return nil if requested.nil?
+
+      s = requested.to_s.downcase
+      return 'none' if s == 'none'
+      return 'gzip' if s == 'gzip'
+      return 'snappy' if s == 'snappy'
+      return 'zstd' if %w[zstd zstandard].include?(s)
+
+      nil
+    end
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: purplelight
 version: !ruby/object:Gem::Version
-  version: 0.1.11
+  version: 0.1.13
 platform: ruby
 authors:
 - Alexander Nicholson