purplelight 0.1.10 → 0.1.12

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f0f51fd601a59915a2a022831663fd4f2468e781b68b96f59d396359be49adbc
-  data.tar.gz: c899a18e7ce390bfc05f832dd32248aa8cbdc7b43bccf86197350e2c7929e7a6
+  metadata.gz: 920b534dc9ac832d83600031277ddd35da2920cff494e0f96d0ca230652d4ba4
+  data.tar.gz: 2fd4476e73efc67d1f4a722dae8d7759ac8b51b5c21546dfb4b06d9fc1cd3934
 SHA512:
-  metadata.gz: f6546911873ed22865b9d4cdd2cc62d855ab3b991030808d8f49f3e054727a406b80c7dc43c518a450915152f2934dcba180d53bf75807c540eef893b3ca50b8
-  data.tar.gz: 5e7176eec64956388e72fd3d894db12e006a18edd4aade7eaf13b144381802932d7207ffefac6ad06157c03363f41acec4ca997871fb8abe8efc9e06e2238804
+  metadata.gz: 370660a815b47c4aa4a0725a6188d6e0455074232000a24c4f909a313c3b9c2d5d3219a17edb628df58f167115b317347c159cc93d5bd59e216d6de1ec7ecd77
+  data.tar.gz: efec787b1e355af50ec07e45b8c125d0f1054640734afbe7790a3a35181f999900d3c843104911eabcea644d2dc7f4604050571e17469f08fe97d9fa5aec92e9

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.10'
+gem 'purplelight', '~> 0.1.12'
 ```
 
 Or install directly:
@@ -144,6 +144,29 @@ Purplelight.snapshot(
 )
 ```
 
+Parquet multi-part (rows-based rotation):
+
+```ruby
+Purplelight.snapshot(
+  client: client,
+  collection: 'users',
+  output: '/data/exports',
+  format: :parquet,
+  # Any mode other than :single_file enables multi-part filenames for Parquet
+  sharding: { mode: :by_size, prefix: 'users_parquet_parts' },
+  # Split into multiple .parquet files, limiting rows per file
+  parquet_max_rows: 100_000,
+  # Optional: Parquet row group size (rows)
+  parquet_row_group: 10_000,
+  resume: { enabled: true }
+)
+```
+
+Notes for Parquet:
+- Parquet multi-part sizing is controlled by rows via `parquet_max_rows`.
+- `--rotate-mb` / `part_bytes` do not affect Parquet part size; they apply to JSONL/CSV.
+- Use `sharding: { mode: :single_file }` to force a single `.parquet` file.
+
 ### Environment variables (optional)
 
 CLI flags take precedence, but these environment variables can set sensible defaults:
@@ -199,6 +222,8 @@ 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.
+- 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
 
@@ -236,6 +261,7 @@ Key points:
 - **Rotation size**: larger (512MB–1GB) reduces finalize overhead for many parts. CLI: `--rotate-mb` (and/or `--by-size`).
 - **JSONL chunking**: tune builder write chunk size for throughput. CLI: `--write-chunk-mb`.
 - **Parquet row groups**: choose a row group size that fits downstream readers. CLI: `--parquet-row-group`.
+- **Parquet parts (rows)**: split Parquet outputs by rows with `parquet_max_rows` (programmatic API). Set `sharding.mode` to anything other than `:single_file` to enable multi-part filenames.
 - **Read preference**: offload to secondaries or tagged analytics nodes when available. CLI: `--read-preference`, `--read-tags`.
 - **Read concern**: pick an appropriate level for consistency/latency trade-offs. CLI: `--read-concern`.
 - **Cursor timeout**: for very long scans, leave `noCursorTimeout` enabled. CLI: `--no-cursor-timeout true|false`.

data/lib/purplelight/snapshot.rb

@@ -39,7 +39,8 @@ module Purplelight
                    sharding: { mode: :by_size, part_bytes: DEFAULTS[:rotate_bytes], prefix: nil },
                    logger: nil, on_progress: nil, read_concern: DEFAULTS[:read_concern], read_preference: DEFAULTS[:read_preference],
                    no_cursor_timeout: DEFAULTS[:no_cursor_timeout], telemetry: nil,
-                   compression_level: nil, writer_threads: 1, write_chunk_bytes: nil, parquet_row_group: nil)
+                   compression_level: nil, writer_threads: 1, write_chunk_bytes: nil, parquet_row_group: nil,
+                   parquet_max_rows: nil)
       @client = client
       @collection = client[collection]
       @output = output
@@ -64,6 +65,7 @@ module Purplelight
       @writer_threads = writer_threads || 1
       @write_chunk_bytes = write_chunk_bytes
       @parquet_row_group = parquet_row_group
+      @parquet_max_rows = parquet_max_rows
 
       @running = true
       @telemetry_enabled = telemetry ? telemetry.enabled? : (ENV['PL_TELEMETRY'] == '1')
@@ -106,6 +108,7 @@ module Purplelight
                             compression_level: @compression_level || (ENV['PL_ZSTD_LEVEL']&.to_i if @compression.to_s == 'zstd') || ENV['PL_ZSTD_LEVEL']&.to_i,
                             write_chunk_bytes: @write_chunk_bytes || ENV['PL_WRITE_CHUNK_BYTES']&.to_i,
                             parquet_row_group: @parquet_row_group || ENV['PL_PARQUET_ROW_GROUP']&.to_i,
+                            parquet_max_rows: @parquet_max_rows,
                             sharding: @sharding,
                             resume_overwrite_incompatible: @resume && @resume[:overwrite_incompatible] ? true : false,
                             telemetry: @telemetry_enabled
@@ -134,7 +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)
+                                   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.10'
+  VERSION = '0.1.12'
 end

data/lib/purplelight/writer_parquet.rb

@@ -15,7 +15,7 @@ module Purplelight
     DEFAULT_ROW_GROUP_SIZE = 10_000
 
     def initialize(directory:, prefix:, compression: :zstd, row_group_size: DEFAULT_ROW_GROUP_SIZE, logger: nil,
-                   manifest: nil, single_file: true, schema: nil)
+                   manifest: nil, single_file: true, schema: nil, rotate_rows: nil)
       @directory = directory
       @prefix = prefix
       @compression = compression
@@ -24,11 +24,13 @@ module Purplelight
       @manifest = manifest
       @single_file = single_file
       @schema = schema
+      @rotate_rows = rotate_rows
 
       @closed = false
       @file_seq = 0
       @part_index = nil
       @pq_writer = nil
+      @rows_in_current_file = 0
 
       ensure_dependencies!
       reset_buffers
@@ -38,15 +40,13 @@ 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
       return if @closed
 
-      ensure_open!
       flush_all_row_groups
-      finalize_current_part!
+      finalize_current_part! if @writer_path
       @closed = true
     end
 
@@ -70,6 +70,7 @@ module Purplelight
       FileUtils.mkdir_p(@directory)
       @writer_path = next_part_path
       @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
@@ -89,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
@@ -101,13 +103,19 @@ 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'
     end
 
     def finalize_current_part!
+      return if @writer_path.nil?
+
       if @pq_writer
         @pq_writer.close
         @pq_writer = nil
@@ -115,6 +123,8 @@ module Purplelight
       @manifest&.complete_part!(index: @part_index, checksum: nil)
       @file_seq += 1 unless @single_file
       @writer_path = nil
+      @part_index = nil
+      @rows_in_current_file = 0
     end
 
     def next_part_path
@@ -147,14 +157,52 @@ module Purplelight
       return if @buffer_docs.empty?
 
       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)
+        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)
+          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
 
@@ -166,14 +214,119 @@ 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)
 
+      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
+      maybe_rotate!
+    end
+
+    def maybe_rotate!
+      return if @single_file
+      return unless @rotate_rows && @rows_in_current_file >= @rotate_rows
+
+      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.10
+  version: 0.1.12
 platform: ruby
 authors:
 - Alexander Nicholson