purplelight 0.1.5 → 0.1.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a650fdd2113129b151396a1a90a83a6f1ede97eb5c34c60e028eb7639d5cc4fd
-  data.tar.gz: 1ab1bc421ddaf1c457639ae2ac5968245c6141b8504897ab762a49bb69c51a69
+  metadata.gz: 6e3771629528ecab067d858491ec4c5de06a8c88c90c64b56dcdb8658c3a6a89
+  data.tar.gz: 4c0427564bf04b5dc3da6b3af97bfd4e5dc0625b2faa0dd8acec84a81bd1c145
 SHA512:
-  metadata.gz: 506e52dce7c474998c8bc4b9afa9f5140349e8e2eed2eed7cccbacac0bd9d9f41528b234f1b99ad8407674791471368ee5b99d93b7ab058522311f2642006a20
-  data.tar.gz: 5c17e387f0d67a21d1351cf4e1e69beaa7beecdf5b9f8011e16bd740e990902abf32c54b02257cdead9c296241557571608b6004446cca5d429675fda07ff61a
+  metadata.gz: 871cf3c04dcaa017d1023ff79135a5050254158a7ba20145abb04c39887125c08eedfb8cae4d16ddd3c4fd0acb383ceb9c11efff943b497ff036b68643dbea79
+  data.tar.gz: 8a5e7fe7a1913939966df61173380e6ed853f45bec2ef90315410a6d36d6341c07d04ff62d200379629b2366cf68296c7c291b5060d0a080dc3015175f1e6780

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.7'
 ```
 
 Or install directly:

data/bin/purplelight

@@ -4,6 +4,7 @@
 require 'optparse'
 require 'json'
 require 'mongo'
+require 'time'
 require_relative '../lib/purplelight'
 
 options = {
@@ -39,7 +40,16 @@ parser = OptionParser.new do |opts|
     options[:sharding] ||= {}
     options[:sharding][:prefix] = v
   end
-  opts.on('-q', '--query JSON', 'Filter query as JSON') { |v| options[:query] = JSON.parse(v) }
+  opts.on('-q', '--query JSON', 'Filter query as JSON (Extended JSON supported)') do |v|
+    begin
+      # Prefer BSON Extended JSON to support $date, $oid, etc.
+      options[:query] = BSON::ExtJSON.parse(v)
+    rescue StandardError
+      # Fallback to plain JSON for compatibility
+      options[:query] = JSON.parse(v)
+    end
+  end
+  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 +116,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/partitioner.rb

@@ -27,13 +27,23 @@ module Purplelight
       max_id = collection.find(query || {}).projection(_id: 1).sort(_id: -1).limit(1).first&.dig('_id')
       return [{ filter: query || {}, sort: { _id: 1 } }] if min_id.nil? || max_id.nil?
 
-      # Create numeric-ish interpolation by sampling
-      ids = collection.find(query || {}).projection(_id: 1).sort(_id: 1).limit(partitions - 1).to_a.map { |d| d['_id'] }
-      boundaries = [min_id] + ids + [max_id]
+      # Create contiguous ranges using ascending inner boundaries.
+      # We intentionally skip the very first _id so the first range includes the smallest document.
+      inner_boundaries = collection.find(query || {})
+                                   .projection(_id: 1)
+                                   .sort(_id: 1)
+                                   .skip(1)
+                                   .limit([partitions - 1, 0].max)
+                                   .to_a
+                                   .map { |d| d['_id'] }
+
       ranges = []
-      boundaries.each_cons(2) do |a, b|
-        ranges << build_range(a, b)
+      prev = nil
+      inner_boundaries.each do |b|
+        ranges << build_range(prev, b)
+        prev = b
       end
+      ranges << build_range(prev, nil)
 
       ranges.map do |r|
         filter = query ? query.dup : {}
@@ -54,7 +64,7 @@ module Purplelight
       min_ts = min_id.respond_to?(:generation_time) ? min_id.generation_time.to_i : nil
       max_ts = max_id.respond_to?(:generation_time) ? max_id.generation_time.to_i : nil
 
-      # Fallback to cursor sampling if _id isn't an ObjectId
+      # Fallback to cursor sampling if _id isn't anObjectId
       return cursor_sampling_partitions(collection: collection, query: query, partitions: partitions) if min_ts.nil? || max_ts.nil? || max_ts <= min_ts
 
       step = [(max_ts - min_ts) / partitions, 1].max

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.7'
 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

data/lib/purplelight/writer_parquet.rb

@@ -116,9 +116,9 @@ module Purplelight
     def next_part_path
       ext = 'parquet'
       filename = if @single_file
-                   format('%<prefix>s.%<ext>s', prefix: @prefix, ext: ext)
+                   "#{@prefix}.#{ext}"
                  else
-                   format('%<prefix}s-part-%<seq>06d.%<ext>s', prefix: @prefix, seq: @file_seq, ext: ext)
+                   format('%<prefix>s-part-%<seq>06d.%<ext>s', prefix: @prefix, seq: @file_seq, ext: ext)
                  end
       File.join(@directory, filename)
     end
@@ -132,7 +132,12 @@ module Purplelight
     end
 
     def extract_value(doc, key)
-      doc[key] || doc[key.to_sym]
+      value = doc[key] || doc[key.to_sym]
+      # Normalize common MongoDB/BSON types to Parquet-friendly values
+      if defined?(BSON) && value.is_a?(BSON::ObjectId)
+        return value.to_s
+      end
+      value
     end
   end
 end

metadata

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