purplelight 0.1.5 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a650fdd2113129b151396a1a90a83a6f1ede97eb5c34c60e028eb7639d5cc4fd
-  data.tar.gz: 1ab1bc421ddaf1c457639ae2ac5968245c6141b8504897ab762a49bb69c51a69
+  metadata.gz: 941ef04ede346a29c0afe19a02d69bd9a25d08ce241e21c6c47960498f4a42c6
+  data.tar.gz: c34b089d2842082f5d4be60c96d58e208e079f92712d397726a964a7361ae895
 SHA512:
-  metadata.gz: 506e52dce7c474998c8bc4b9afa9f5140349e8e2eed2eed7cccbacac0bd9d9f41528b234f1b99ad8407674791471368ee5b99d93b7ab058522311f2642006a20
-  data.tar.gz: 5c17e387f0d67a21d1351cf4e1e69beaa7beecdf5b9f8011e16bd740e990902abf32c54b02257cdead9c296241557571608b6004446cca5d429675fda07ff61a
+  metadata.gz: 8d666cb565283e6410fb0412d8ed369db18b0eaa0e80c573d677295910319b809396e99fca6627214d15942d531c9ac378c7d5a850842a7124d2ff4c9b03d7d8
+  data.tar.gz: 9188cd0b55e0d3e54fb2d209b85eeeacfb1c0d8daa73b1ace0fbb3bf4a2af865d6dd5c125f1e487458f096077a70cc93e7c0191c1222f6edd58701487f5479da

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.5'
+gem 'purplelight', '~> 0.1.6'
 ```
 
 Or install directly:

data/bin/purplelight

@@ -40,6 +40,7 @@ parser = OptionParser.new do |opts|
     options[:sharding][:prefix] = v
   end
   opts.on('-q', '--query JSON', 'Filter query as JSON') { |v| options[:query] = JSON.parse(v) }
+  opts.on('--projection JSON', 'Projection as JSON (e.g., {"_id":1,"field":1})') { |v| options[:projection] = JSON.parse(v) }
   opts.on('--read-preference MODE',
           'Read preference mode: primary|primary_preferred|secondary|secondary_preferred|nearest') do |v|
     options[:read_preference] = v.to_sym
@@ -106,6 +107,7 @@ ok = Purplelight.snapshot(
   partitions: options[:partitions],
   batch_size: options[:batch_size],
   query: options[:query],
+  projection: options[:projection],
   sharding: options[:sharding],
   read_preference: effective_read || options[:read_preference],
   resume: { enabled: true },

data/lib/purplelight/snapshot.rb

@@ -221,6 +221,8 @@ module Purplelight
       cursor = @collection.find(filter, opts)
 
       encode_lines = (@format == :jsonl)
+      # When JSONL, build one big string per batch to offload join cost from writer.
+      string_batch = +''
       buffer = []
       buffer_bytes = 0
       last_id = checkpoint
@@ -230,27 +232,43 @@ module Purplelight
           doc = @mapper.call(doc) if @mapper
           t_ser = telemetry.start(:serialize_time)
           if encode_lines
-            line = "#{JSON.generate(doc)}\n"
+            line = "#{JSON.fast_generate(doc)}\n"
             telemetry.finish(:serialize_time, t_ser)
             bytes = line.bytesize
-            buffer << line
+            string_batch << line
           else
             # For CSV/Parquet keep raw docs to allow schema/row building
-            bytes = (JSON.generate(doc).bytesize + 1)
+            bytes = (JSON.fast_generate(doc).bytesize + 1)
             telemetry.finish(:serialize_time, t_ser)
             buffer << doc
           end
           buffer_bytes += bytes
-          next unless buffer.length >= batch_size || buffer_bytes >= 1_000_000
+          # For JSONL, we count rows via newline accumulation; for others, use array length
+          ready = encode_lines ? (buffer_bytes >= 1_000_000 || (string_batch.length >= 1_000_000)) : (buffer.length >= batch_size || buffer_bytes >= 1_000_000)
+          next unless ready
 
           t_q = telemetry.start(:queue_wait_time)
-          queue.push(buffer, bytes: buffer_bytes)
+          if encode_lines
+            queue.push(string_batch, bytes: buffer_bytes)
+            string_batch = +''
+          else
+            queue.push(buffer, bytes: buffer_bytes)
+            buffer = []
+          end
           telemetry.finish(:queue_wait_time, t_q)
           manifest.update_partition_checkpoint!(idx, last_id)
-          buffer = []
           buffer_bytes = 0
         end
-        unless buffer.empty?
+        if encode_lines
+          unless string_batch.empty?
+            t_q = telemetry.start(:queue_wait_time)
+            queue.push(string_batch, bytes: buffer_bytes)
+            telemetry.finish(:queue_wait_time, t_q)
+            manifest.update_partition_checkpoint!(idx, last_id)
+            string_batch = +''
+            buffer_bytes = 0
+          end
+        elsif !buffer.empty?
           t_q = telemetry.start(:queue_wait_time)
           queue.push(buffer, bytes: buffer_bytes)
           telemetry.finish(:queue_wait_time, t_q)

data/lib/purplelight/version.rb

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

data/lib/purplelight/writer_csv.rb

@@ -11,6 +11,12 @@ rescue LoadError
   # zstd not available; fallback handled later via gzip
 end
 
+begin
+  require 'zstd-ruby'
+rescue LoadError
+  # alternative zstd gem not available
+end
+
 module Purplelight
   # WriterCSV writes documents to CSV files with optional compression.
   class WriterCSV
@@ -24,6 +30,8 @@ module Purplelight
       @rotate_bytes = rotate_bytes
       @logger = logger
       @manifest = manifest
+      env_level = ENV['PL_ZSTD_LEVEL']&.to_i
+      @compression_level = (env_level && env_level > 0 ? env_level : nil)
       @single_file = single_file
 
       @columns = columns&.map(&:to_s)
@@ -91,6 +99,39 @@ module Purplelight
 
     private
 
+    # Minimal wrapper to count bytes written for rotate logic when
+    # underlying compressed writer doesn't expose position (e.g., zstd-ruby).
+    class CountingIO
+      def initialize(io, on_write:)
+        @io = io
+        @on_write = on_write
+      end
+
+      def write(data)
+        bytes_written = @io.write(data)
+        @on_write.call(bytes_written) if bytes_written && @on_write
+        bytes_written
+      end
+
+      # CSV calls '<<' on the underlying IO in some code paths
+      def <<(data)
+        write(data)
+      end
+
+      # CSV#flush may forward flush to underlying IO; make it a no-op if unavailable
+      def flush
+        @io.flush if @io.respond_to?(:flush)
+      end
+
+      def method_missing(method_name, *args, &block)
+        @io.send(method_name, *args, &block)
+      end
+
+      def respond_to_missing?(method_name, include_private = false)
+        @io.respond_to?(method_name, include_private)
+      end
+    end
+
     def ensure_open!
       return if @io
 
@@ -98,7 +139,8 @@ module Purplelight
       path = next_part_path
       @part_index = @manifest&.open_part!(path) if @manifest
       raw = File.open(path, 'wb')
-      @io = build_compressed_io(raw)
+      compressed = build_compressed_io(raw)
+      @io = CountingIO.new(compressed, on_write: ->(n) { @bytes_written += n })
       @csv = CSV.new(@io)
       @bytes_written = 0
       @rows_written = 0
@@ -107,7 +149,13 @@ module Purplelight
     def build_compressed_io(raw)
       case @effective_compression.to_s
       when 'zstd'
-        return ZSTDS::Writer.open(raw, level: 10) if defined?(ZSTDS)
+        if Object.const_defined?(:Zstd) && defined?(::Zstd::StreamWriter)
+          level = @compression_level || 10
+          return ::Zstd::StreamWriter.new(raw, level: level)
+        elsif defined?(ZSTDS)
+          level = @compression_level || 10
+          return ZSTDS::Stream::Writer.new(raw, compression_level: level)
+        end
 
         @logger&.warn('zstd gem not loaded; using gzip')
         Zlib::GzipWriter.new(raw)
@@ -154,7 +202,7 @@ module Purplelight
     def determine_effective_compression(requested)
       case requested.to_s
       when 'zstd'
-        (defined?(ZSTDS) ? :zstd : :gzip)
+        ((defined?(ZSTDS) || (Object.const_defined?(:Zstd) && defined?(::Zstd::StreamWriter))) ? :zstd : :gzip)
       when 'none'
         :none
       else

data/lib/purplelight/writer_jsonl.rb

@@ -10,6 +10,12 @@ rescue LoadError
   # zstd not available; will fallback to gzip
 end
 
+begin
+  require 'zstd-ruby'
+rescue LoadError
+  # alternative zstd gem not available
+end
+
 module Purplelight
   # WriterJSONL writes newline-delimited JSON with optional compression.
   class WriterJSONL
@@ -23,7 +29,8 @@ module Purplelight
       @rotate_bytes = rotate_bytes
       @logger = logger
       @manifest = manifest
-      @compression_level = compression_level
+      env_level = ENV['PL_ZSTD_LEVEL']&.to_i
+      @compression_level = compression_level || (env_level && env_level > 0 ? env_level : nil)
 
       @part_index = nil
       @io = nil
@@ -33,23 +40,71 @@ module Purplelight
       @closed = false
 
       @effective_compression = determine_effective_compression(@compression)
-      return unless @effective_compression.to_s != @compression.to_s
-
-      @logger&.warn("requested compression '#{@compression}' not available; using '#{@effective_compression}'")
+      if @logger
+        level_disp = @compression_level || (ENV['PL_ZSTD_LEVEL']&.to_i if @effective_compression.to_s == 'zstd')
+        @logger.info("WriterJSONL using compression='#{@effective_compression}' level='#{level_disp || 'default'}'")
+      end
+      if @effective_compression.to_s != @compression.to_s
+        @logger&.warn("requested compression '#{@compression}' not available; using '#{@effective_compression}'")
+      end
     end
 
-    def write_many(array_of_docs)
+    def write_many(batch)
       ensure_open!
-      # If upstream already produced newline-terminated strings, join fast.
-      buffer = if array_of_docs.first.is_a?(String)
-                 array_of_docs.join
-               else
-                 array_of_docs.map { |doc| "#{JSON.generate(doc)}\n" }.join
-               end
-      rows = array_of_docs.size
-      write_buffer(buffer)
+
+      chunk_threshold = (ENV['PL_WRITE_CHUNK_BYTES']&.to_i || (8 * 1024 * 1024))
+      total_bytes = 0
+      rows = 0
+
+      if batch.is_a?(String)
+        # Fast-path: writer received a preassembled buffer string
+        buffer = batch
+        rows = buffer.count("\n")
+        write_buffer(buffer)
+        total_bytes = buffer.bytesize
+      elsif batch.first.is_a?(String)
+        # Join and write in chunks to avoid large intermediate allocations
+        chunk = +''
+        chunk_bytes = 0
+        batch.each do |line|
+          chunk << line
+          rows += 1
+          chunk_bytes += line.bytesize
+          next unless chunk_bytes >= chunk_threshold
+
+          write_buffer(chunk)
+          total_bytes += chunk.bytesize
+          chunk = +''
+          chunk_bytes = 0
+        end
+        unless chunk.empty?
+          write_buffer(chunk)
+          total_bytes += chunk.bytesize
+        end
+      else
+        # Fallback: encode docs here (JSON.fast_generate preferred) and write in chunks
+        chunk = +''
+        chunk_bytes = 0
+        batch.each do |doc|
+          line = "#{JSON.fast_generate(doc)}\n"
+          rows += 1
+          chunk << line
+          chunk_bytes += line.bytesize
+          next unless chunk_bytes >= chunk_threshold
+
+          write_buffer(chunk)
+          total_bytes += chunk.bytesize
+          chunk = +''
+          chunk_bytes = 0
+        end
+        unless chunk.empty?
+          write_buffer(chunk)
+          total_bytes += chunk.bytesize
+        end
+      end
+
       @rows_written += rows
-      @manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows, bytes_delta: buffer.bytesize)
+      @manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows, bytes_delta: total_bytes)
     end
 
     def rotate_if_needed
@@ -86,15 +141,18 @@ module Purplelight
     def build_compressed_io(raw)
       case @effective_compression.to_s
       when 'zstd'
-        if defined?(ZSTDS)
-          # ZSTDS::Writer supports IO-like interface
+        # Prefer zstd-ruby if available, else ruby-zstds
+        if Object.const_defined?(:Zstd) && defined?(::Zstd::StreamWriter)
+          level = @compression_level || 3
+          return ::Zstd::StreamWriter.new(raw, level: level)
+        elsif defined?(ZSTDS)
           level = @compression_level || 3
-          ZSTDS::Writer.open(raw, level: level)
-        else
-          @logger&.warn('zstd gem not loaded; this should have been handled earlier')
-          level = @compression_level || Zlib::DEFAULT_COMPRESSION
-          Zlib::GzipWriter.new(raw, level)
+          return ZSTDS::Stream::Writer.new(raw, compression_level: level)
         end
+
+        @logger&.warn('zstd gems not loaded; falling back to gzip')
+        level = @compression_level || Zlib::DEFAULT_COMPRESSION
+        Zlib::GzipWriter.new(raw, level)
       when 'gzip'
         level = @compression_level || 1
         Zlib::GzipWriter.new(raw, level)
@@ -142,7 +200,7 @@ module Purplelight
     def determine_effective_compression(requested)
       case requested.to_s
       when 'zstd'
-        (defined?(ZSTDS) ? :zstd : :gzip)
+        ((defined?(ZSTDS) || (Object.const_defined?(:Zstd) && defined?(::Zstd::StreamWriter))) ? :zstd : :gzip)
       when 'none'
         :none
       else

metadata

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