purplelight 0.1.13 → 0.1.17

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,344 @@
1
+ #!/usr/bin/env ruby
2
+ # frozen_string_literal: true
3
+
4
+ require 'bundler/setup'
5
+ require 'fileutils'
6
+ require 'json'
7
+ require 'stringio'
8
+ require 'tmpdir'
9
+ require 'time'
10
+ require 'purplelight'
11
+ require_relative 'harness'
12
+
13
+ DOCS = Array.new(200) do |index|
14
+ {
15
+ '_id' => BSON::ObjectId.new,
16
+ 'index' => index,
17
+ 'active' => index.even?,
18
+ 'tags' => %w[alpha beta],
19
+ 'nested' => { 'value' => index }
20
+ }
21
+ end.freeze
22
+ PARQUET_DOCS = Array.new(10_000) do |index|
23
+ {
24
+ '_id' => BSON::ObjectId.new,
25
+ 'index' => index,
26
+ 'active' => index.even?,
27
+ 'tags' => %w[alpha beta],
28
+ 'nested' => { 'value' => index }
29
+ }
30
+ end.freeze
31
+ ENCODED_JSONL = DOCS.map { |document| "#{JSON.generate(document)}\n" }.join * 10
32
+ ENCODED_BATCH = Purplelight::WriterJSONL::EncodedBatch.new(
33
+ data: ENCODED_JSONL,
34
+ rows: DOCS.length * 10,
35
+ bytes: ENCODED_JSONL.bytesize
36
+ )
37
+ PARTITION_DOCS = Array.new(5_001) { |index| { '_id' => index } }.freeze
38
+ pipeline_random = Random.new(12_345)
39
+ PIPELINE_DOCS = Array.new(20_000) do |index|
40
+ object_id = BSON::ObjectId.new.tap(&:to_s)
41
+ { '_id' => object_id, 'index' => index, 'payload' => pipeline_random.bytes(512).unpack1('H*') }
42
+ end.freeze
43
+
44
+ FakeCollection = Data.define(:name)
45
+ FakeClient = Class.new do
46
+ def [](name)
47
+ FakeCollection.new(name.to_s)
48
+ end
49
+ end
50
+ EmptyCursor = Class.new do
51
+ def projection(*) = self
52
+ def sort(*) = self
53
+ def limit(*) = self
54
+ def first = nil
55
+ def each = nil
56
+ end
57
+
58
+ EmptyCollection = Data.define(:name) do
59
+ def find(*) = EmptyCursor.new
60
+ def estimated_document_count = 0
61
+ end
62
+
63
+ EmptyClient = Class.new do
64
+ def [](name) = EmptyCollection.new(name.to_s)
65
+ end
66
+
67
+ RawCursor = Class.new do
68
+ def initialize(documents)
69
+ @documents = documents
70
+ end
71
+
72
+ def each(&) = @documents.each(&)
73
+ end
74
+
75
+ RawCollection = Data.define(:name, :documents) do
76
+ def find(*) = RawCursor.new(documents)
77
+ end
78
+
79
+ CollectionClient = Data.define(:collection) do
80
+ def [](*) = collection
81
+ end
82
+
83
+ PipelineCursor = Class.new do
84
+ include Enumerable
85
+
86
+ def initialize(documents)
87
+ @documents = documents
88
+ @direction = 1
89
+ @limit = nil
90
+ end
91
+
92
+ def projection(*) = self
93
+ def hint(*) = self
94
+ def batch_size(*) = self
95
+ def no_cursor_timeout = self
96
+ def skip(*) = self
97
+
98
+ def sort(specification)
99
+ @direction = specification[:_id] || specification['_id'] || 1
100
+ self
101
+ end
102
+
103
+ def limit(count)
104
+ @limit = count
105
+ self
106
+ end
107
+
108
+ def to_a
109
+ documents = @direction.negative? ? @documents.reverse : @documents
110
+ @limit ? documents.first(@limit) : documents
111
+ end
112
+
113
+ def first = to_a.first
114
+ def each(&) = to_a.each(&)
115
+ end
116
+
117
+ PipelineCollection = Data.define(:name, :documents) do
118
+ def find(*) = PipelineCursor.new(documents)
119
+ def estimated_document_count = documents.length
120
+ end
121
+
122
+ SinkQueue = Class.new do
123
+ def push(*) = nil
124
+ end
125
+
126
+ SinkManifest = Class.new do
127
+ def partition_checkpoint(*) = nil
128
+ def update_partition_checkpoint!(*) = nil
129
+ def mark_partition_complete!(*) = nil
130
+ end
131
+
132
+ PlannerCursor = Class.new do
133
+ include Enumerable
134
+
135
+ def initialize(documents)
136
+ @documents = documents
137
+ end
138
+
139
+ def projection(*) = self
140
+ def batch_size(*) = self
141
+ def no_cursor_timeout = self
142
+ def each(&) = @documents.each(&)
143
+ end
144
+
145
+ PlannerCollection = Data.define(:documents) do
146
+ def find(*) = PlannerCursor.new(documents)
147
+ def estimated_document_count = 10_001
148
+ end
149
+
150
+ harness = Purplelight::Microbench::Harness.new
151
+
152
+ Dir.mktmpdir('purplelight-microbench') do |directory|
153
+ sequence = 0
154
+
155
+ harness.register('byte_queue_round_trip', paths: :queue, iterations: 20_000) do
156
+ queue = Purplelight::ByteQueue.new(max_bytes: 1024)
157
+ queue.push(sequence, bytes: 8)
158
+ queue.pop
159
+ end
160
+ harness.register('byte_queue_bulk_drain', paths: :queue, iterations: 10) do
161
+ queue = Purplelight::ByteQueue.new(max_bytes: 50_000)
162
+ 5_000.times { queue.push(sequence, bytes: 8) }
163
+ 5_000.times { queue.pop }
164
+ end
165
+
166
+ telemetry = Purplelight::Telemetry.new
167
+ harness.register('telemetry_counter_and_timer', paths: :telemetry, iterations: 50_000) do
168
+ ticket = telemetry.start(:operation)
169
+ telemetry.add(:documents)
170
+ telemetry.finish(:operation, ticket)
171
+ end
172
+
173
+ harness.register('telemetry_construction', paths: :telemetry, iterations: 100_000) do
174
+ Purplelight::Telemetry.new
175
+ end
176
+
177
+ harness.register('manifest_progress', paths: :manifest, iterations: 250) do
178
+ sequence += 1
179
+ path = File.join(directory, "manifest-#{sequence}.json")
180
+ manifest = Purplelight::Manifest.new(path:)
181
+ manifest.configure!(collection: 'records', format: :jsonl, compression: :none,
182
+ query_digest: Purplelight::Manifest.query_digest({}, nil))
183
+ manifest.ensure_partitions!(1)
184
+ index = manifest.open_part!("part-#{sequence}")
185
+ manifest.add_progress_to_part!(index:, rows_delta: 200, bytes_delta: 4096)
186
+ manifest.update_partition_checkpoint!(0, sequence)
187
+ manifest.mark_partition_complete!(0)
188
+ manifest.complete_part!(index:)
189
+ end
190
+
191
+ harness.register('partition_range_planning', paths: :partitioner, iterations: 100_000) do
192
+ Purplelight::Partitioner.build_range(sequence, sequence + 1)
193
+ end
194
+
195
+ planner_collection = PlannerCollection.new(PARTITION_DOCS)
196
+ harness.register('partition_cursor_sampling', paths: :partitioner, iterations: 10) do
197
+ Purplelight::Partitioner.cursor_sampling_partitions(collection: planner_collection, query: nil, partitions: 2)
198
+ end
199
+
200
+ client = FakeClient.new
201
+ harness.register('snapshot_configuration', paths: :snapshot, iterations: 10_000) do
202
+ snapshot = Purplelight::Snapshot.new(client:, collection: :records, output: directory,
203
+ format: :jsonl, compression: :none, partitions: 1)
204
+ snapshot.send(:resolve_output, directory, :jsonl)
205
+ end
206
+ empty_client = EmptyClient.new
207
+ harness.register('snapshot_empty_export', paths: :snapshot, iterations: 20) do
208
+ sequence += 1
209
+ output = File.join(directory, "snapshot-#{sequence}")
210
+ Purplelight::Snapshot.new(client: empty_client, collection: :records, output:, format: :jsonl,
211
+ compression: :none, partitions: 1, resume: { enabled: false }).run
212
+ end
213
+
214
+ raw_collection = RawCollection.new('records', DOCS)
215
+ raw_snapshot = Purplelight::Snapshot.new(client: CollectionClient.new(raw_collection), collection: :records,
216
+ output: directory, format: :csv, compression: :none, batch_size: DOCS.length,
217
+ partitions: 1, resume: { enabled: false })
218
+ sink_queue = SinkQueue.new
219
+ sink_manifest = SinkManifest.new
220
+ harness.register('snapshot_raw_batching', paths: :snapshot, iterations: 100) do
221
+ raw_snapshot.send(:read_partition, idx: 0, filter_spec: { filter: {} }, queue: sink_queue,
222
+ batch_size: DOCS.length, manifest: sink_manifest)
223
+ end
224
+
225
+ pipeline_collection = PipelineCollection.new('pipeline', PIPELINE_DOCS)
226
+ pipeline_client = CollectionClient.new(pipeline_collection)
227
+ harness.register('snapshot_jsonl_serial_pipeline', paths: %i[manifest queue snapshot writer_jsonl], iterations: 3) do
228
+ sequence += 1
229
+ output = File.join(directory, "pipeline-serial-#{sequence}.jsonl")
230
+ Purplelight::Snapshot.new(client: pipeline_client, collection: :pipeline, output:, format: :jsonl,
231
+ compression: :zstd, batch_size: 250, partitions: 1, writer_threads: 1,
232
+ resume: { enabled: false }).run
233
+ end
234
+
235
+ harness.register('snapshot_jsonl_pipeline', paths: %i[manifest queue snapshot writer_jsonl], iterations: 3) do
236
+ sequence += 1
237
+ output = File.join(directory, "pipeline-#{sequence}.jsonl")
238
+ Purplelight::Snapshot.new(client: pipeline_client, collection: :pipeline, output:, format: :jsonl,
239
+ compression: :zstd, batch_size: 250, partitions: 1, writer_threads: 2,
240
+ resume: { enabled: false }).run
241
+ end
242
+
243
+ harness.register('jsonl_batch_write', paths: :writer_jsonl, iterations: 100) do
244
+ sequence += 1
245
+ writer = Purplelight::WriterJSONL.new(directory:, prefix: "jsonl-#{sequence}", compression: :none)
246
+ writer.write_many(DOCS)
247
+ writer.close
248
+ end
249
+
250
+ harness.register('jsonl_zstd_batch_write', paths: :writer_jsonl, iterations: 50) do
251
+ sequence += 1
252
+ writer = Purplelight::WriterJSONL.new(directory:, prefix: "jsonl-zstd-#{sequence}", compression: :zstd)
253
+ writer.write_many(DOCS)
254
+ writer.close
255
+ end
256
+
257
+ accounting_writer = Purplelight::WriterJSONL.new(
258
+ directory:, prefix: 'accounting', compression: :none, rotate_bytes: nil
259
+ )
260
+ accounting_io = StringIO.new
261
+ accounting_writer.instance_variable_set(:@io, accounting_io)
262
+ harness.register('jsonl_preencoded_accounting', paths: :writer_jsonl, iterations: 10_000) do
263
+ accounting_io.truncate(0)
264
+ accounting_io.rewind
265
+ accounting_writer.write_many(ENCODED_BATCH)
266
+ end
267
+
268
+ harness.register('csv_batch_write', paths: :writer_csv, iterations: 100) do
269
+ sequence += 1
270
+ writer = Purplelight::WriterCSV.new(directory:, prefix: "csv-#{sequence}", compression: :none,
271
+ single_file: true)
272
+ writer.write_many(DOCS)
273
+ writer.close
274
+ end
275
+
276
+ harness.register('csv_zstd_batch_write', paths: :writer_csv, iterations: 50) do
277
+ sequence += 1
278
+ writer = Purplelight::WriterCSV.new(directory:, prefix: "csv-zstd-#{sequence}", compression: :zstd,
279
+ single_file: true)
280
+ writer.write_many(DOCS)
281
+ writer.close
282
+ end
283
+
284
+ harness.register('csv_gzip_batch_write', paths: :writer_csv, iterations: 50) do
285
+ sequence += 1
286
+ writer = Purplelight::WriterCSV.new(directory:, prefix: "csv-gzip-#{sequence}", compression: :gzip,
287
+ compression_level: 1, single_file: true)
288
+ writer.write_many(DOCS)
289
+ writer.close
290
+ end
291
+ raw_csv_io = StringIO.new
292
+ counting_io_class = Purplelight::WriterCSV::CountingIO
293
+ has_callback = counting_io_class.instance_method(:initialize).parameters.any? { |_, name| name == :on_write }
294
+ counting_io = if has_callback
295
+ counting_io_class.new(raw_csv_io, on_write: ->(*) {})
296
+ else
297
+ counting_io_class.new(raw_csv_io)
298
+ end
299
+ harness.register('csv_counting_io_write', paths: :writer_csv, iterations: 100_000) do
300
+ counting_io.write('payload')
301
+ end
302
+ csv_serializer = Purplelight::WriterCSV.new(directory:, prefix: 'csv-serializer', compression: :none,
303
+ columns: DOCS.first.keys)
304
+ csv_serialized_row = +''
305
+ harness.register('csv_row_serialization', paths: :writer_csv, iterations: 25_000) do
306
+ csv_serialized_row.clear
307
+ csv_serializer.send(:append_csv_document, csv_serialized_row, DOCS.first)
308
+ end
309
+
310
+ parquet_buffer_writer = Purplelight::WriterParquet.new(directory:, prefix: 'parquet-buffer', compression: :none,
311
+ row_group_size: 1_000_000, single_file: true)
312
+ harness.register('parquet_buffer_append', paths: :writer_parquet, iterations: 5_000) do
313
+ parquet_buffer_writer.write_many(DOCS)
314
+ parquet_buffer_writer.instance_variable_get(:@buffer_docs).clear
315
+ end
316
+
317
+ object_ids = DOCS.map { |document| document.fetch('_id') }.freeze
318
+ harness.register('parquet_object_id_array', paths: :writer_parquet, iterations: 1_000) do
319
+ parquet_buffer_writer.send(:build_object_id_array, object_ids)
320
+ end
321
+
322
+ harness.register('parquet_batch_write', paths: :writer_parquet, iterations: 10) do
323
+ sequence += 1
324
+ writer = Purplelight::WriterParquet.new(directory:, prefix: "parquet-#{sequence}", compression: :zstd,
325
+ row_group_size: 100, single_file: true)
326
+ writer.write_many(DOCS)
327
+ writer.close
328
+ end
329
+
330
+ harness.register('parquet_row_group_write', paths: :writer_parquet, iterations: 2) do
331
+ sequence += 1
332
+ writer = Purplelight::WriterParquet.new(directory:, prefix: "parquet-row-groups-#{sequence}",
333
+ compression: :zstd, row_group_size: 1_000, single_file: true)
334
+ writer.write_many(PARQUET_DOCS)
335
+ writer.close
336
+ end
337
+
338
+ baseline_index = ARGV.index('--baseline')
339
+ output_index = ARGV.index('--output')
340
+ report = harness.run(baseline_path: baseline_index && ARGV.fetch(baseline_index + 1))
341
+ rendered = JSON.pretty_generate(report)
342
+ puts rendered
343
+ File.write(ARGV.fetch(output_index + 1), "#{rendered}\n") if output_index
344
+ end
@@ -0,0 +1,33 @@
1
+ # frozen_string_literal: true
2
+
3
+ require 'bundler/setup'
4
+ require 'fileutils'
5
+ require 'optparse'
6
+ require 'stackprof'
7
+
8
+ mode = :cpu
9
+ OptionParser.new do |parser|
10
+ parser.on('--mode MODE', %w[cpu object], 'StackProf sampling mode') { |value| mode = value.to_sym }
11
+ end.parse!
12
+
13
+ root = File.expand_path('..', __dir__)
14
+ profile_directory = File.join(root, 'tmp', 'profiles')
15
+ FileUtils.mkdir_p(profile_directory)
16
+ dump_path = File.join(profile_directory, "#{mode}.dump")
17
+ flamegraph_path = File.join(profile_directory, "#{mode}-flamegraph.html")
18
+ benchmark_output = File.join(profile_directory, "#{mode}-benchmark.json")
19
+ interval = mode == :object ? 100 : 1_000
20
+
21
+ original_arguments = ARGV.dup
22
+ ARGV.replace(['--output', benchmark_output])
23
+ StackProf.run(mode:, interval:, raw: true, out: dump_path) do
24
+ load File.join(__dir__, 'microbench.rb')
25
+ end
26
+ ARGV.replace(original_arguments)
27
+
28
+ File.open(flamegraph_path, 'w') do |file|
29
+ StackProf::Report.from_file(dump_path).print_d3_flamegraph(file)
30
+ end
31
+
32
+ puts "StackProf dump: #{dump_path}"
33
+ puts "Flamegraph: #{flamegraph_path}"
data/bin/purplelight CHANGED
@@ -59,7 +59,7 @@ parser = OptionParser.new do |opts|
59
59
  options[:sharding] ||= {}
60
60
  options[:sharding][:prefix] = v
61
61
  end
62
- opts.on('--writer-threads N', Integer, 'Number of writer threads (experimental, JSONL only)') { |v| options[:writer_threads] = v }
62
+ opts.on('--writer-threads N', Integer, 'Number of concurrent JSONL writer threads') { |v| options[:writer_threads] = v }
63
63
  opts.on('--write-chunk-mb MB', Integer, 'JSONL encode/write chunk size in MB') { |v| options[:write_chunk_bytes] = v * 1024 * 1024 }
64
64
  opts.on('--parquet-row-group N', Integer, 'Parquet row group size (rows)') { |v| options[:parquet_row_group] = v }
65
65
  opts.on('--parquet-max-rows N', Integer, 'Parquet max rows per part (multi-part only)') { |v| options[:parquet_max_rows] = v }
@@ -160,8 +160,8 @@ snapshot_args[:no_cursor_timeout] = options[:no_cursor_timeout] unless options[:
160
160
  snapshot_args[:compression_level] = options[:compression_level] if options[:compression_level]
161
161
  snapshot_args[:writer_threads] = options[:writer_threads] if options[:writer_threads]
162
162
  snapshot_args[:write_chunk_bytes] = options[:write_chunk_bytes] if options[:write_chunk_bytes]
163
- snapshot_args[:parquet_row_group] = options[:parquet_row_group] if options[:parquet_row_group]
164
- snapshot_args[:parquet_max_rows] = options[:parquet_max_rows] if options[:parquet_max_rows]
163
+ snapshot_args[:parquet_row_group] = options[:parquet_row_group] if options[:parquet_row_group]
164
+ snapshot_args[:parquet_max_rows] = options[:parquet_max_rows] if options[:parquet_max_rows]
165
165
 
166
166
  # telemetry env override
167
167
  if options[:telemetry_flag]
@@ -172,7 +172,7 @@ end
172
172
  ENV['PL_ZSTD_LEVEL'] = options[:compression_level].to_s if options[:compression_level]
173
173
  ENV['PL_WRITE_CHUNK_BYTES'] = options[:write_chunk_bytes].to_s if options[:write_chunk_bytes]
174
174
  ENV['PL_PARQUET_ROW_GROUP'] = options[:parquet_row_group].to_s if options[:parquet_row_group]
175
- # No env default for parquet_max_rows; it is snapshot-argument only
175
+ # No env default for parquet_max_rows; it is snapshot-argument only
176
176
  ENV['PL_WRITER_THREADS'] = options[:writer_threads].to_s if options[:writer_threads]
177
177
 
178
178
  ok = Purplelight.snapshot(**snapshot_args)
@@ -1,5 +1,6 @@
1
1
  # frozen_string_literal: true
2
2
 
3
+ require 'bson'
3
4
  require 'json'
4
5
  require 'time'
5
6
  require 'securerandom'
@@ -13,7 +14,7 @@ module Purplelight
13
14
  # counts so interrupted runs can resume safely and completed runs are
14
15
  # reproducible. Methods are thread-safe where mutation occurs.
15
16
  class Manifest
16
- DEFAULT_VERSION = 1
17
+ DEFAULT_VERSION = 2
17
18
 
18
19
  attr_reader :path, :data
19
20
 
@@ -24,6 +25,9 @@ module Purplelight
24
25
 
25
26
  def initialize(path:, data: nil)
26
27
  @path = path
28
+ @directory = File.dirname(path)
29
+ @temporary_path = "#{path}.tmp"
30
+ FileUtils.mkdir_p(@directory)
27
31
  @data = data || {
28
32
  'version' => DEFAULT_VERSION,
29
33
  'run_id' => SecureRandom.uuid,
@@ -33,11 +37,12 @@ module Purplelight
33
37
  'compression' => nil,
34
38
  'query_digest' => nil,
35
39
  'options' => {},
40
+ 'partition_filters' => nil,
36
41
  'parts' => [],
37
42
  'partitions' => []
38
43
  }
39
44
  @mutex = Mutex.new
40
- @last_save_at = Time.now
45
+ @last_save_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
41
46
  end
42
47
 
43
48
  def self.load(path)
@@ -46,23 +51,25 @@ module Purplelight
46
51
  end
47
52
 
48
53
  def save!
49
- dir = File.dirname(path)
50
- FileUtils.mkdir_p(dir)
51
- tmp = "#{path}.tmp"
52
- File.write(tmp, JSON.pretty_generate(@data))
53
- FileUtils.mv(tmp, path)
54
- end
55
-
56
- def configure!(collection:, format:, compression:, query_digest:, options: {})
57
- @data['collection'] = collection
58
- @data['format'] = format.to_s
59
- @data['compression'] = compression.to_s
60
- @data['query_digest'] = query_digest
61
- @data['options'] = options
62
- save!
54
+ File.write(@temporary_path, JSON.pretty_generate(@data))
55
+ File.rename(@temporary_path, path)
56
+ end
57
+
58
+ def configure!(collection:, format:, compression:, query_digest:, options: {}, partition_count: nil)
59
+ @mutex.synchronize do
60
+ @data['collection'] = collection
61
+ @data['format'] = format.to_s
62
+ @data['compression'] = compression.to_s
63
+ @data['query_digest'] = query_digest
64
+ @data['options'] = options
65
+ normalize_partition_data!(partition_count) if partition_count
66
+ save!
67
+ end
63
68
  end
64
69
 
65
70
  def compatible_with?(collection:, format:, compression:, query_digest:)
71
+ return false unless @data['version'] == DEFAULT_VERSION
72
+
66
73
  @data['collection'] == collection &&
67
74
  @data['format'] == format.to_s &&
68
75
  @data['compression'] == compression.to_s &&
@@ -72,9 +79,7 @@ module Purplelight
72
79
  def ensure_partitions!(count)
73
80
  @mutex.synchronize do
74
81
  if @data['partitions'].empty?
75
- @data['partitions'] = Array.new(count) do |i|
76
- { 'index' => i, 'last_id_exclusive' => nil, 'completed' => false }
77
- end
82
+ normalize_partition_data!(count)
78
83
  save!
79
84
  end
80
85
  end
@@ -83,8 +88,8 @@ module Purplelight
83
88
  def update_partition_checkpoint!(index, last_id_exclusive)
84
89
  @mutex.synchronize do
85
90
  part = @data['partitions'][index]
86
- part['last_id_exclusive'] = last_id_exclusive
87
- save!
91
+ part['last_id_exclusive'] = serialize_checkpoint(last_id_exclusive)
92
+ save_maybe!
88
93
  end
89
94
  end
90
95
 
@@ -132,10 +137,51 @@ module Purplelight
132
137
  @data['partitions']
133
138
  end
134
139
 
140
+ def partition_checkpoint(index)
141
+ raw_checkpoint = @data['partitions'][index]&.fetch('last_id_exclusive', nil)
142
+ BSON::ExtJSON.parse_obj(raw_checkpoint) if raw_checkpoint
143
+ end
144
+
145
+ def partition_filters
146
+ raw_filters = @data['partition_filters']
147
+ return unless raw_filters
148
+
149
+ BSON::ExtJSON.parse_obj(raw_filters).map do |filter_spec|
150
+ {
151
+ filter: filter_spec[:filter] || filter_spec['filter'],
152
+ sort: filter_spec[:sort] || filter_spec['sort'],
153
+ hint: filter_spec[:hint] || filter_spec['hint']
154
+ }
155
+ end
156
+ end
157
+
158
+ def configure_partition_filters!(filters)
159
+ @mutex.synchronize do
160
+ @data['partition_filters'] = filters.as_extended_json
161
+ end
162
+ end
163
+
135
164
  private
136
165
 
166
+ def serialize_checkpoint(value)
167
+ case value
168
+ when nil, String, Numeric, true, false
169
+ value
170
+ else
171
+ value.as_extended_json
172
+ end
173
+ end
174
+
175
+ def normalize_partition_data!(count)
176
+ return unless @data['partitions'].empty?
177
+
178
+ @data['partitions'] = Array.new(count) do |index|
179
+ { 'index' => index, 'last_id_exclusive' => nil, 'completed' => false }
180
+ end
181
+ end
182
+
137
183
  def save_maybe!(interval_seconds: 2.0)
138
- now = Time.now
184
+ now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
139
185
  return unless (now - @last_save_at) >= interval_seconds
140
186
 
141
187
  save!
@@ -1,6 +1,7 @@
1
1
  # frozen_string_literal: true
2
2
 
3
3
  require 'mongo'
4
+ require_relative 'telemetry'
4
5
 
5
6
  module Purplelight
6
7
  # Partitioner builds MongoDB range filters to split work across workers.
@@ -14,7 +15,7 @@ module Purplelight
14
15
  def self.object_id_partitions(collection:, query:, partitions:, mode: nil, telemetry: nil)
15
16
  # Choose planning mode: :timestamp (fast), :cursor (legacy)
16
17
  chosen_mode = (mode || ENV['PL_PARTITIONER_MODE'] || :timestamp).to_sym
17
- telemetry ||= (defined?(Telemetry) ? Telemetry::NULL : nil)
18
+ telemetry ||= Telemetry::NULL
18
19
 
19
20
  return cursor_sampling_partitions(collection: collection, query: query, partitions: partitions) if chosen_mode == :cursor
20
21
 
@@ -112,11 +113,16 @@ module Purplelight
112
113
  step = [total / partitions, 1].max
113
114
  boundaries = []
114
115
  cursor = base_query.projection(_id: 1).batch_size(1_000).no_cursor_timeout
115
- i = 0
116
+ index = 0
117
+ next_boundary = 0
116
118
  cursor.each do |doc|
117
- boundaries << doc['_id'] if (i % step).zero?
118
- i += 1
119
- break if boundaries.size >= partitions
119
+ if index == next_boundary
120
+ boundaries << doc['_id']
121
+ break if boundaries.size >= partitions
122
+
123
+ next_boundary += step
124
+ end
125
+ index += 1
120
126
  end
121
127
 
122
128
  ranges = []
@@ -6,6 +6,7 @@ module Purplelight
6
6
  def initialize(max_bytes: 128 * 1024 * 1024)
7
7
  @max_bytes = max_bytes
8
8
  @queue = []
9
+ @sizes = []
9
10
  @bytes = 0
10
11
  @closed = false
11
12
  @mutex = Mutex.new
@@ -16,8 +17,12 @@ module Purplelight
16
17
  @mutex.synchronize do
17
18
  raise 'queue closed' if @closed
18
19
 
19
- @cv.wait(@mutex) while (@bytes + bytes) > @max_bytes
20
- @queue << [item, bytes]
20
+ while !@queue.empty? && (@bytes + bytes) > @max_bytes
21
+ @cv.wait(@mutex)
22
+ raise 'queue closed' if @closed
23
+ end
24
+ @queue << item
25
+ @sizes << bytes
21
26
  @bytes += bytes
22
27
  @cv.broadcast
23
28
  end
@@ -30,7 +35,8 @@ module Purplelight
30
35
 
31
36
  @cv.wait(@mutex)
32
37
  end
33
- item, bytes = @queue.shift
38
+ item = @queue.shift
39
+ bytes = @sizes.shift
34
40
  @bytes -= bytes
35
41
  @cv.broadcast
36
42
  item