purplelight 0.1.13 → 0.1.16
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.
- checksums.yaml +4 -4
- data/PERFORMANCE_NOTES.md +53 -0
- data/README.md +19 -5
- data/Rakefile +11 -0
- data/benchmark/baseline.json +366 -0
- data/benchmark/harness.rb +109 -0
- data/benchmark/microbench.rb +344 -0
- data/benchmark/profile.rb +33 -0
- data/bin/purplelight +4 -4
- data/lib/purplelight/manifest.rb +68 -22
- data/lib/purplelight/partitioner.rb +11 -5
- data/lib/purplelight/queue.rb +9 -3
- data/lib/purplelight/snapshot.rb +91 -55
- data/lib/purplelight/telemetry.rb +0 -1
- data/lib/purplelight/version.rb +1 -1
- data/lib/purplelight/writer_csv.rb +76 -45
- data/lib/purplelight/writer_jsonl.rb +57 -24
- data/lib/purplelight/writer_parquet.rb +268 -145
- metadata +14 -15
data/lib/purplelight/snapshot.rb
CHANGED
|
@@ -19,7 +19,7 @@ module Purplelight
|
|
|
19
19
|
format: :jsonl,
|
|
20
20
|
compression: :zstd,
|
|
21
21
|
batch_size: 2_000,
|
|
22
|
-
partitions:
|
|
22
|
+
partitions: (Etc.nprocessors * 2).clamp(4, 32),
|
|
23
23
|
queue_size_bytes: 256 * 1024 * 1024,
|
|
24
24
|
rotate_bytes: 256 * 1024 * 1024,
|
|
25
25
|
read_concern: { level: :majority },
|
|
@@ -68,6 +68,10 @@ module Purplelight
|
|
|
68
68
|
@parquet_max_rows = parquet_max_rows
|
|
69
69
|
|
|
70
70
|
@running = true
|
|
71
|
+
if @on_progress
|
|
72
|
+
@progress_mutex = Mutex.new
|
|
73
|
+
@progress_cv = ConditionVariable.new
|
|
74
|
+
end
|
|
71
75
|
@telemetry_enabled = telemetry ? telemetry.enabled? : (ENV['PL_TELEMETRY'] == '1')
|
|
72
76
|
@telemetry = telemetry || (
|
|
73
77
|
@telemetry_enabled ? Telemetry.new(enabled: true) : Telemetry::NULL
|
|
@@ -105,43 +109,55 @@ module Purplelight
|
|
|
105
109
|
read_concern: (@read_concern.is_a?(Hash) ? @read_concern : { level: @read_concern }),
|
|
106
110
|
no_cursor_timeout: @no_cursor_timeout,
|
|
107
111
|
writer_threads: @writer_threads,
|
|
108
|
-
compression_level: @compression_level ||
|
|
112
|
+
compression_level: @compression_level || ENV['PL_ZSTD_LEVEL']&.to_i,
|
|
109
113
|
write_chunk_bytes: @write_chunk_bytes || ENV['PL_WRITE_CHUNK_BYTES']&.to_i,
|
|
110
114
|
parquet_row_group: @parquet_row_group || ENV['PL_PARQUET_ROW_GROUP']&.to_i,
|
|
111
115
|
parquet_max_rows: @parquet_max_rows,
|
|
112
116
|
sharding: @sharding,
|
|
113
117
|
resume_overwrite_incompatible: @resume && @resume[:overwrite_incompatible] ? true : false,
|
|
114
118
|
telemetry: @telemetry_enabled
|
|
115
|
-
})
|
|
116
|
-
manifest.ensure_partitions!(@partitions)
|
|
119
|
+
}, partition_count: @partitions)
|
|
117
120
|
|
|
118
|
-
#
|
|
121
|
+
# Reuse the original ranges on resume. Replanning against a growing collection
|
|
122
|
+
# can move partition boundaries and replay rows under the old checkpoints.
|
|
119
123
|
t_plan = @telemetry.start(:partition_plan_time)
|
|
120
|
-
partition_filters =
|
|
121
|
-
|
|
124
|
+
partition_filters = manifest.partition_filters
|
|
125
|
+
unless partition_filters
|
|
126
|
+
partition_filters = Partitioner.object_id_partitions(collection: @collection, query: @query,
|
|
127
|
+
partitions: @partitions, telemetry: @telemetry)
|
|
128
|
+
manifest.configure_partition_filters!(partition_filters)
|
|
129
|
+
end
|
|
122
130
|
@telemetry.finish(:partition_plan_time, t_plan)
|
|
123
131
|
|
|
124
132
|
# Reader queue
|
|
125
133
|
queue = ByteQueue.new(max_bytes: @queue_size_bytes)
|
|
126
134
|
|
|
127
135
|
# Writer
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
136
|
+
writer_count = @format == :jsonl ? [@writer_threads.to_i, 1].max : 1
|
|
137
|
+
part_sequence = WriterJSONL::PartSequence.new(manifest.parts.length) if @format == :jsonl
|
|
138
|
+
writers = case @format
|
|
139
|
+
when :jsonl
|
|
140
|
+
Array.new(writer_count) do
|
|
141
|
+
WriterJSONL.new(directory: dir, prefix: prefix, compression: @compression,
|
|
142
|
+
rotate_bytes: @rotate_bytes, logger: @logger, manifest: manifest,
|
|
143
|
+
compression_level: @compression_level, write_chunk_bytes: @write_chunk_bytes,
|
|
144
|
+
part_sequence:)
|
|
145
|
+
end
|
|
146
|
+
when :csv
|
|
147
|
+
single_file = @sharding && @sharding[:mode].to_s == 'single_file'
|
|
148
|
+
[WriterCSV.new(directory: dir, prefix: prefix, compression: @compression,
|
|
149
|
+
rotate_bytes: @rotate_bytes, logger: @logger, manifest: manifest,
|
|
150
|
+
single_file:, compression_level: @compression_level)]
|
|
151
|
+
when :parquet
|
|
152
|
+
single_file = @sharding && @sharding[:mode].to_s == 'single_file'
|
|
153
|
+
row_group = @parquet_row_group || ENV['PL_PARQUET_ROW_GROUP']&.to_i ||
|
|
154
|
+
WriterParquet::DEFAULT_ROW_GROUP_SIZE
|
|
155
|
+
[WriterParquet.new(directory: dir, prefix: prefix, compression: @compression, logger: @logger,
|
|
156
|
+
manifest:, single_file:, row_group_size: row_group,
|
|
157
|
+
rotate_rows: @parquet_max_rows)]
|
|
158
|
+
else
|
|
159
|
+
raise ArgumentError, "format not implemented: #{@format}"
|
|
160
|
+
end
|
|
145
161
|
|
|
146
162
|
# Start reader threads
|
|
147
163
|
readers = partition_filters.each_with_index.map do |pf, idx|
|
|
@@ -153,37 +169,49 @@ module Purplelight
|
|
|
153
169
|
end
|
|
154
170
|
end
|
|
155
171
|
|
|
156
|
-
# Writer
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
172
|
+
# Writer loops
|
|
173
|
+
writer_telemetries = []
|
|
174
|
+
writer_threads = writers.map do |writer|
|
|
175
|
+
writer_telemetry = @telemetry_enabled ? Telemetry.new(enabled: true) : Telemetry::NULL
|
|
176
|
+
writer_telemetries << writer_telemetry
|
|
177
|
+
Thread.new(writer, writer_telemetry) do |worker_writer, worker_telemetry|
|
|
178
|
+
Thread.current[:pl_telemetry] = worker_telemetry if @telemetry_enabled
|
|
179
|
+
loop do
|
|
180
|
+
batch = queue.pop
|
|
181
|
+
break if batch.nil?
|
|
163
182
|
|
|
164
|
-
|
|
183
|
+
worker_writer.write_many(batch)
|
|
184
|
+
end
|
|
185
|
+
ensure
|
|
186
|
+
worker_writer.close
|
|
165
187
|
end
|
|
166
|
-
ensure
|
|
167
|
-
writer.close
|
|
168
188
|
end
|
|
169
189
|
|
|
170
|
-
progress_thread =
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
190
|
+
progress_thread = if @on_progress
|
|
191
|
+
Thread.new do
|
|
192
|
+
loop do
|
|
193
|
+
@progress_mutex.synchronize { @progress_cv.wait(@progress_mutex, 2) }
|
|
194
|
+
break unless @running
|
|
175
195
|
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
196
|
+
@on_progress.call({ queue_bytes: queue.size_bytes })
|
|
197
|
+
end
|
|
198
|
+
end
|
|
199
|
+
end
|
|
179
200
|
|
|
180
201
|
# Join readers
|
|
181
202
|
readers.each(&:join)
|
|
182
203
|
queue.close
|
|
183
|
-
|
|
184
|
-
@telemetry.merge!(writer_telemetry) if @telemetry_enabled
|
|
185
|
-
|
|
186
|
-
|
|
204
|
+
writer_threads.each(&:join)
|
|
205
|
+
writer_telemetries.each { |writer_telemetry| @telemetry.merge!(writer_telemetry) } if @telemetry_enabled
|
|
206
|
+
if progress_thread
|
|
207
|
+
@progress_mutex.synchronize do
|
|
208
|
+
@running = false
|
|
209
|
+
@progress_cv.broadcast
|
|
210
|
+
end
|
|
211
|
+
progress_thread.join
|
|
212
|
+
else
|
|
213
|
+
@running = false
|
|
214
|
+
end
|
|
187
215
|
if @telemetry_enabled
|
|
188
216
|
total = @telemetry.timers.values.sum
|
|
189
217
|
breakdown = @telemetry.timers
|
|
@@ -222,7 +250,7 @@ module Purplelight
|
|
|
222
250
|
hint = @hint || filter_spec[:hint] || { _id: 1 }
|
|
223
251
|
|
|
224
252
|
# Resume from checkpoint if present
|
|
225
|
-
checkpoint = manifest.
|
|
253
|
+
checkpoint = manifest.partition_checkpoint(idx)
|
|
226
254
|
if checkpoint
|
|
227
255
|
filter = filter.dup
|
|
228
256
|
filter['_id'] = (filter['_id'] || {}).merge({ '$gt' => checkpoint })
|
|
@@ -248,11 +276,12 @@ module Purplelight
|
|
|
248
276
|
string_batch = +''
|
|
249
277
|
buffer = []
|
|
250
278
|
buffer_bytes = 0
|
|
279
|
+
buffer_rows = 0
|
|
251
280
|
json_state = if encode_lines
|
|
252
281
|
JSON::Ext::Generator::State.new(ascii_only: false, max_nesting: false,
|
|
253
282
|
buffer_initial_length: 4_096)
|
|
254
283
|
end
|
|
255
|
-
|
|
284
|
+
bson_size_buffer = BSON::ByteBuffer.new unless encode_lines
|
|
256
285
|
last_id = checkpoint
|
|
257
286
|
begin
|
|
258
287
|
cursor.each do |doc|
|
|
@@ -266,37 +295,43 @@ module Purplelight
|
|
|
266
295
|
string_batch << "\n"
|
|
267
296
|
bytes = json.bytesize + 1
|
|
268
297
|
else
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
bytes =
|
|
298
|
+
previous_length = bson_size_buffer.length
|
|
299
|
+
doc.to_bson(bson_size_buffer)
|
|
300
|
+
bytes = bson_size_buffer.length - previous_length
|
|
272
301
|
telemetry.finish(:serialize_time, t_ser)
|
|
273
302
|
buffer << doc
|
|
274
303
|
end
|
|
275
304
|
buffer_bytes += bytes
|
|
276
|
-
|
|
277
|
-
|
|
305
|
+
buffer_rows += 1
|
|
306
|
+
# Flush on the configured row count or before a batch exceeds 1 MB.
|
|
307
|
+
ready = buffer_rows >= batch_size || buffer_bytes >= 1_000_000
|
|
278
308
|
next unless ready
|
|
279
309
|
|
|
280
310
|
t_q = telemetry.start(:queue_wait_time)
|
|
281
311
|
if encode_lines
|
|
282
|
-
|
|
312
|
+
batch = WriterJSONL::EncodedBatch.new(data: string_batch, rows: buffer_rows, bytes: buffer_bytes)
|
|
313
|
+
queue.push(batch, bytes: buffer_bytes)
|
|
283
314
|
string_batch = +''
|
|
284
315
|
else
|
|
285
316
|
queue.push(buffer, bytes: buffer_bytes)
|
|
286
317
|
buffer = []
|
|
318
|
+
bson_size_buffer = BSON::ByteBuffer.new
|
|
287
319
|
end
|
|
288
320
|
telemetry.finish(:queue_wait_time, t_q)
|
|
289
321
|
manifest.update_partition_checkpoint!(idx, last_id)
|
|
290
322
|
buffer_bytes = 0
|
|
323
|
+
buffer_rows = 0
|
|
291
324
|
end
|
|
292
325
|
if encode_lines
|
|
293
326
|
unless string_batch.empty?
|
|
294
327
|
t_q = telemetry.start(:queue_wait_time)
|
|
295
|
-
|
|
328
|
+
batch = WriterJSONL::EncodedBatch.new(data: string_batch, rows: buffer_rows, bytes: buffer_bytes)
|
|
329
|
+
queue.push(batch, bytes: buffer_bytes)
|
|
296
330
|
telemetry.finish(:queue_wait_time, t_q)
|
|
297
331
|
manifest.update_partition_checkpoint!(idx, last_id)
|
|
298
332
|
string_batch = +''
|
|
299
333
|
buffer_bytes = 0
|
|
334
|
+
buffer_rows = 0
|
|
300
335
|
end
|
|
301
336
|
elsif !buffer.empty?
|
|
302
337
|
t_q = telemetry.start(:queue_wait_time)
|
|
@@ -305,6 +340,7 @@ module Purplelight
|
|
|
305
340
|
manifest.update_partition_checkpoint!(idx, last_id)
|
|
306
341
|
buffer = []
|
|
307
342
|
buffer_bytes = 0
|
|
343
|
+
buffer_rows = 0
|
|
308
344
|
end
|
|
309
345
|
manifest.mark_partition_complete!(idx)
|
|
310
346
|
end
|
data/lib/purplelight/version.rb
CHANGED
|
@@ -1,10 +1,10 @@
|
|
|
1
1
|
# frozen_string_literal: true
|
|
2
2
|
|
|
3
|
-
require 'csv'
|
|
4
3
|
require 'json'
|
|
5
4
|
require 'zlib'
|
|
6
5
|
require 'fileutils'
|
|
7
6
|
|
|
7
|
+
# simplecov:disable
|
|
8
8
|
begin
|
|
9
9
|
require 'zstd-ruby'
|
|
10
10
|
rescue LoadError
|
|
@@ -14,14 +14,17 @@ rescue LoadError
|
|
|
14
14
|
# no zstd backend; gzip fallback used later
|
|
15
15
|
end
|
|
16
16
|
end
|
|
17
|
+
# simplecov:enable
|
|
17
18
|
|
|
18
19
|
module Purplelight
|
|
19
20
|
# WriterCSV writes documents to CSV files with optional compression.
|
|
20
21
|
class WriterCSV
|
|
21
22
|
DEFAULT_ROTATE_BYTES = 256 * 1024 * 1024
|
|
23
|
+
DEFAULT_ZSTD_LEVEL = 9
|
|
24
|
+
DEFAULT_GZIP_LEVEL = 1
|
|
22
25
|
|
|
23
26
|
def initialize(directory:, prefix:, compression: :zstd, rotate_bytes: DEFAULT_ROTATE_BYTES, logger: nil,
|
|
24
|
-
manifest: nil, single_file: false, columns: nil, headers: true)
|
|
27
|
+
manifest: nil, single_file: false, columns: nil, headers: true, compression_level: nil)
|
|
25
28
|
@directory = directory
|
|
26
29
|
@prefix = prefix
|
|
27
30
|
@compression = compression
|
|
@@ -29,7 +32,7 @@ module Purplelight
|
|
|
29
32
|
@logger = logger
|
|
30
33
|
@manifest = manifest
|
|
31
34
|
env_level = ENV['PL_ZSTD_LEVEL']&.to_i
|
|
32
|
-
@compression_level = (env_level&.positive? ? env_level : nil)
|
|
35
|
+
@compression_level = compression_level || (env_level&.positive? ? env_level : nil)
|
|
33
36
|
@single_file = single_file
|
|
34
37
|
|
|
35
38
|
@columns = columns&.map(&:to_s)
|
|
@@ -37,10 +40,7 @@ module Purplelight
|
|
|
37
40
|
|
|
38
41
|
@part_index = nil
|
|
39
42
|
@io = nil
|
|
40
|
-
@
|
|
41
|
-
@bytes_written = 0
|
|
42
|
-
@rows_written = 0
|
|
43
|
-
@file_seq = 0
|
|
43
|
+
@file_seq = manifest ? manifest.parts.length : 0
|
|
44
44
|
@closed = false
|
|
45
45
|
|
|
46
46
|
@effective_compression = determine_effective_compression(@compression)
|
|
@@ -52,22 +52,24 @@ module Purplelight
|
|
|
52
52
|
def write_many(array_of_docs)
|
|
53
53
|
ensure_open!
|
|
54
54
|
|
|
55
|
-
|
|
55
|
+
output = +''
|
|
56
|
+
# Infer columns if needed from docs.
|
|
56
57
|
if @columns.nil?
|
|
57
58
|
sample_docs = array_of_docs.is_a?(Array) ? array_of_docs : []
|
|
58
|
-
sample_docs = sample_docs.
|
|
59
|
+
sample_docs = sample_docs.grep_v(String)
|
|
59
60
|
@columns = infer_columns(sample_docs)
|
|
60
|
-
|
|
61
|
+
append_csv_row(output, @columns) if @headers
|
|
61
62
|
end
|
|
62
63
|
|
|
64
|
+
rows = 0
|
|
63
65
|
array_of_docs.each do |doc|
|
|
64
66
|
next if doc.is_a?(String)
|
|
65
67
|
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
@rows_written += 1
|
|
68
|
+
append_csv_document(output, doc)
|
|
69
|
+
rows += 1
|
|
69
70
|
end
|
|
70
|
-
@
|
|
71
|
+
@io.write(output) unless output.empty?
|
|
72
|
+
@manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows, bytes_delta: 0)
|
|
71
73
|
|
|
72
74
|
rotate_if_needed
|
|
73
75
|
end
|
|
@@ -76,7 +78,7 @@ module Purplelight
|
|
|
76
78
|
return if @single_file
|
|
77
79
|
return if @rotate_bytes.nil?
|
|
78
80
|
|
|
79
|
-
raw_bytes = @io.respond_to?(:pos) ? @io.pos : @bytes_written
|
|
81
|
+
raw_bytes = @io.respond_to?(:pos) ? @io.pos : @io.bytes_written
|
|
80
82
|
return if raw_bytes < @rotate_bytes
|
|
81
83
|
|
|
82
84
|
rotate!
|
|
@@ -85,7 +87,6 @@ module Purplelight
|
|
|
85
87
|
def close
|
|
86
88
|
return if @closed
|
|
87
89
|
|
|
88
|
-
@csv&.flush
|
|
89
90
|
if @io
|
|
90
91
|
t = Thread.current[:pl_telemetry]&.start(:rotate_time)
|
|
91
92
|
finalize_current_part!
|
|
@@ -100,26 +101,20 @@ module Purplelight
|
|
|
100
101
|
# Minimal wrapper to count bytes written for rotate logic when
|
|
101
102
|
# underlying compressed writer doesn't expose position (e.g., zstd-ruby).
|
|
102
103
|
class CountingIO
|
|
103
|
-
|
|
104
|
+
attr_reader :bytes_written
|
|
105
|
+
|
|
106
|
+
def initialize(io)
|
|
104
107
|
@io = io
|
|
105
|
-
@
|
|
108
|
+
@bytes_written = 0
|
|
106
109
|
end
|
|
107
110
|
|
|
108
111
|
def write(data)
|
|
109
112
|
bytes_written = @io.write(data)
|
|
110
|
-
@
|
|
113
|
+
@bytes_written += bytes_written
|
|
111
114
|
bytes_written
|
|
112
115
|
end
|
|
113
116
|
|
|
114
|
-
|
|
115
|
-
def <<(data)
|
|
116
|
-
write(data)
|
|
117
|
-
end
|
|
118
|
-
|
|
119
|
-
# CSV#flush may forward flush to underlying IO; make it a no-op if unavailable
|
|
120
|
-
def flush
|
|
121
|
-
@io.flush if @io.respond_to?(:flush)
|
|
122
|
-
end
|
|
117
|
+
alias << write
|
|
123
118
|
|
|
124
119
|
def method_missing(method_name, *, &)
|
|
125
120
|
@io.send(method_name, *, &)
|
|
@@ -135,31 +130,33 @@ module Purplelight
|
|
|
135
130
|
|
|
136
131
|
FileUtils.mkdir_p(@directory)
|
|
137
132
|
path = next_part_path
|
|
138
|
-
@part_index = @manifest
|
|
133
|
+
@part_index = @manifest.open_part!(path) if @manifest
|
|
134
|
+
# The compression stream owns and closes this file handle.
|
|
135
|
+
# rubocop:disable Style/FileOpen
|
|
139
136
|
raw = File.open(path, 'wb')
|
|
137
|
+
# rubocop:enable Style/FileOpen
|
|
140
138
|
compressed = build_compressed_io(raw)
|
|
141
|
-
@io = CountingIO.new(compressed
|
|
142
|
-
@
|
|
143
|
-
|
|
144
|
-
|
|
139
|
+
@io = CountingIO.new(compressed)
|
|
140
|
+
return unless @headers && @columns
|
|
141
|
+
|
|
142
|
+
header = +''
|
|
143
|
+
append_csv_row(header, @columns)
|
|
144
|
+
@io.write(header)
|
|
145
145
|
end
|
|
146
146
|
|
|
147
147
|
def build_compressed_io(raw)
|
|
148
148
|
case @effective_compression.to_s
|
|
149
149
|
when 'zstd'
|
|
150
|
+
level = @compression_level || DEFAULT_ZSTD_LEVEL
|
|
150
151
|
if Object.const_defined?(:Zstd) && defined?(::Zstd::StreamWriter)
|
|
151
|
-
level
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
level = @compression_level || 10
|
|
155
|
-
return ZSTDS::Stream::Writer.new(raw, compression_level: level)
|
|
152
|
+
::Zstd::StreamWriter.new(raw, level: level)
|
|
153
|
+
else
|
|
154
|
+
ZSTDS::Stream::Writer.new(raw, compression_level: level)
|
|
156
155
|
end
|
|
157
156
|
|
|
158
|
-
@logger&.warn('zstd gem not loaded; using gzip')
|
|
159
|
-
Zlib::GzipWriter.new(raw)
|
|
160
|
-
|
|
161
157
|
when 'gzip'
|
|
162
|
-
|
|
158
|
+
level = @compression_level || DEFAULT_GZIP_LEVEL
|
|
159
|
+
Zlib::GzipWriter.new(raw, level)
|
|
163
160
|
when 'none'
|
|
164
161
|
raw
|
|
165
162
|
else
|
|
@@ -175,7 +172,6 @@ module Purplelight
|
|
|
175
172
|
@io.close
|
|
176
173
|
Thread.current[:pl_telemetry]&.finish(:rotate_time, t)
|
|
177
174
|
@io = nil
|
|
178
|
-
@csv = nil
|
|
179
175
|
ensure_open!
|
|
180
176
|
end
|
|
181
177
|
|
|
@@ -216,6 +212,39 @@ module Purplelight
|
|
|
216
212
|
:gzip
|
|
217
213
|
end
|
|
218
214
|
|
|
215
|
+
def append_csv_document(output, document)
|
|
216
|
+
index = 0
|
|
217
|
+
while index < @columns.length
|
|
218
|
+
output << ',' unless index.zero?
|
|
219
|
+
append_csv_value(output, extract_value(document, @columns[index]))
|
|
220
|
+
index += 1
|
|
221
|
+
end
|
|
222
|
+
output << "\n"
|
|
223
|
+
end
|
|
224
|
+
|
|
225
|
+
def append_csv_row(output, values)
|
|
226
|
+
values.each_with_index do |value, index|
|
|
227
|
+
output << ',' unless index.zero?
|
|
228
|
+
append_csv_value(output, value)
|
|
229
|
+
end
|
|
230
|
+
output << "\n"
|
|
231
|
+
end
|
|
232
|
+
|
|
233
|
+
def append_csv_value(output, value)
|
|
234
|
+
return if value.nil?
|
|
235
|
+
|
|
236
|
+
string = value.to_s
|
|
237
|
+
contains_quote = string.include?('"')
|
|
238
|
+
unless string.empty? || contains_quote || string.include?(',') || string.include?("\n") || string.include?("\r")
|
|
239
|
+
output << string
|
|
240
|
+
return
|
|
241
|
+
end
|
|
242
|
+
|
|
243
|
+
output << '"'
|
|
244
|
+
output << (contains_quote ? string.gsub('"', '""') : string)
|
|
245
|
+
output << '"'
|
|
246
|
+
end
|
|
247
|
+
|
|
219
248
|
def infer_columns(docs)
|
|
220
249
|
keys = {}
|
|
221
250
|
docs.each do |d|
|
|
@@ -223,12 +252,14 @@ module Purplelight
|
|
|
223
252
|
end
|
|
224
253
|
# Put _id first if present, then other keys sorted
|
|
225
254
|
cols = []
|
|
226
|
-
|
|
255
|
+
first = docs.first
|
|
256
|
+
cols << '_id' if first && (first.key?('_id') || first.key?(:_id))
|
|
227
257
|
cols + keys.keys.sort
|
|
228
258
|
end
|
|
229
259
|
|
|
230
260
|
def extract_value(doc, key)
|
|
231
|
-
val = doc[key]
|
|
261
|
+
val = doc[key]
|
|
262
|
+
val = doc[key.to_sym] if val.nil? && !doc.key?(key)
|
|
232
263
|
case val
|
|
233
264
|
when Hash, Array
|
|
234
265
|
JSON.generate(val)
|
|
@@ -4,6 +4,7 @@ require 'json'
|
|
|
4
4
|
require 'zlib'
|
|
5
5
|
require 'fileutils'
|
|
6
6
|
|
|
7
|
+
# simplecov:disable
|
|
7
8
|
begin
|
|
8
9
|
require 'zstd-ruby'
|
|
9
10
|
rescue LoadError
|
|
@@ -13,14 +14,34 @@ rescue LoadError
|
|
|
13
14
|
# no zstd backend; gzip fallback
|
|
14
15
|
end
|
|
15
16
|
end
|
|
17
|
+
# simplecov:enable
|
|
16
18
|
|
|
17
19
|
module Purplelight
|
|
18
20
|
# WriterJSONL writes newline-delimited JSON with optional compression.
|
|
19
21
|
class WriterJSONL
|
|
20
22
|
DEFAULT_ROTATE_BYTES = 256 * 1024 * 1024
|
|
23
|
+
DEFAULT_ZSTD_LEVEL = 3
|
|
24
|
+
DEFAULT_GZIP_LEVEL = 1
|
|
25
|
+
EncodedBatch = Data.define(:data, :rows, :bytes)
|
|
26
|
+
|
|
27
|
+
# Allocates globally unique output part numbers across writer threads.
|
|
28
|
+
class PartSequence
|
|
29
|
+
def initialize(next_value = 0)
|
|
30
|
+
@next_value = next_value
|
|
31
|
+
@mutex = Mutex.new
|
|
32
|
+
end
|
|
33
|
+
|
|
34
|
+
def next
|
|
35
|
+
@mutex.synchronize do
|
|
36
|
+
value = @next_value
|
|
37
|
+
@next_value += 1
|
|
38
|
+
value
|
|
39
|
+
end
|
|
40
|
+
end
|
|
41
|
+
end
|
|
21
42
|
|
|
22
43
|
def initialize(directory:, prefix:, compression: :zstd, rotate_bytes: DEFAULT_ROTATE_BYTES, logger: nil,
|
|
23
|
-
manifest: nil, compression_level: nil)
|
|
44
|
+
manifest: nil, compression_level: nil, write_chunk_bytes: nil, part_sequence: nil)
|
|
24
45
|
@directory = directory
|
|
25
46
|
@prefix = prefix
|
|
26
47
|
@compression = compression
|
|
@@ -29,18 +50,20 @@ module Purplelight
|
|
|
29
50
|
@manifest = manifest
|
|
30
51
|
env_level = ENV['PL_ZSTD_LEVEL']&.to_i
|
|
31
52
|
@compression_level = compression_level || (env_level&.positive? ? env_level : nil)
|
|
53
|
+
@write_chunk_bytes = write_chunk_bytes
|
|
54
|
+
@part_sequence = part_sequence
|
|
32
55
|
|
|
33
56
|
@part_index = nil
|
|
34
57
|
@io = nil
|
|
35
58
|
@bytes_written = 0
|
|
36
|
-
@
|
|
37
|
-
@file_seq = 0
|
|
59
|
+
@file_seq = manifest ? manifest.parts.length : 0
|
|
38
60
|
@closed = false
|
|
61
|
+
@thread_telemetry = false
|
|
39
62
|
|
|
40
63
|
@effective_compression = determine_effective_compression(@compression)
|
|
41
64
|
@json_state = JSON::Ext::Generator::State.new(ascii_only: false, max_nesting: false)
|
|
42
65
|
if @logger
|
|
43
|
-
level_disp = @compression_level
|
|
66
|
+
level_disp = @compression_level
|
|
44
67
|
@logger.info("WriterJSONL using compression='#{@effective_compression}' level='#{level_disp || 'default'}'")
|
|
45
68
|
end
|
|
46
69
|
return unless @effective_compression.to_s != @compression.to_s
|
|
@@ -51,12 +74,16 @@ module Purplelight
|
|
|
51
74
|
def write_many(batch)
|
|
52
75
|
ensure_open!
|
|
53
76
|
|
|
54
|
-
chunk_threshold = ENV['PL_WRITE_CHUNK_BYTES']&.to_i || (8 * 1024 * 1024)
|
|
77
|
+
chunk_threshold = @write_chunk_bytes || ENV['PL_WRITE_CHUNK_BYTES']&.to_i || (8 * 1024 * 1024)
|
|
55
78
|
total_bytes = 0
|
|
56
79
|
rows = 0
|
|
57
80
|
|
|
58
|
-
if batch.is_a?(
|
|
59
|
-
|
|
81
|
+
if batch.is_a?(EncodedBatch)
|
|
82
|
+
write_buffer(batch.data)
|
|
83
|
+
rows = batch.rows
|
|
84
|
+
total_bytes = batch.bytes
|
|
85
|
+
elsif batch.is_a?(String)
|
|
86
|
+
# Fast path for callers that don't provide row metadata.
|
|
60
87
|
buffer = batch
|
|
61
88
|
rows = buffer.count("\n")
|
|
62
89
|
write_buffer(buffer)
|
|
@@ -104,7 +131,6 @@ module Purplelight
|
|
|
104
131
|
end
|
|
105
132
|
end
|
|
106
133
|
|
|
107
|
-
@rows_written += rows
|
|
108
134
|
@manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows, bytes_delta: total_bytes)
|
|
109
135
|
end
|
|
110
136
|
|
|
@@ -132,30 +158,27 @@ module Purplelight
|
|
|
132
158
|
|
|
133
159
|
FileUtils.mkdir_p(@directory)
|
|
134
160
|
path = next_part_path
|
|
135
|
-
@part_index = @manifest
|
|
161
|
+
@part_index = @manifest.open_part!(path) if @manifest
|
|
162
|
+
# The compression stream owns and closes this file handle.
|
|
163
|
+
# rubocop:disable Style/FileOpen
|
|
136
164
|
raw = File.open(path, 'wb')
|
|
165
|
+
# rubocop:enable Style/FileOpen
|
|
137
166
|
@io = build_compressed_io(raw)
|
|
138
167
|
@bytes_written = 0
|
|
139
|
-
@rows_written = 0
|
|
140
168
|
end
|
|
141
169
|
|
|
142
170
|
def build_compressed_io(raw)
|
|
143
171
|
case @effective_compression.to_s
|
|
144
172
|
when 'zstd'
|
|
145
173
|
# Prefer zstd-ruby if available, else ruby-zstds
|
|
174
|
+
level = @compression_level || DEFAULT_ZSTD_LEVEL
|
|
146
175
|
if Object.const_defined?(:Zstd) && defined?(::Zstd::StreamWriter)
|
|
147
|
-
level
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
level = @compression_level || 3
|
|
151
|
-
return ZSTDS::Stream::Writer.new(raw, compression_level: level)
|
|
176
|
+
::Zstd::StreamWriter.new(raw, level: level)
|
|
177
|
+
else
|
|
178
|
+
ZSTDS::Stream::Writer.new(raw, compression_level: level)
|
|
152
179
|
end
|
|
153
|
-
|
|
154
|
-
@logger&.warn('zstd gems not loaded; falling back to gzip')
|
|
155
|
-
level = @compression_level || Zlib::DEFAULT_COMPRESSION
|
|
156
|
-
Zlib::GzipWriter.new(raw, level)
|
|
157
180
|
when 'gzip'
|
|
158
|
-
level = @compression_level ||
|
|
181
|
+
level = @compression_level || DEFAULT_GZIP_LEVEL
|
|
159
182
|
Zlib::GzipWriter.new(raw, level)
|
|
160
183
|
when 'none'
|
|
161
184
|
raw
|
|
@@ -165,9 +188,18 @@ module Purplelight
|
|
|
165
188
|
end
|
|
166
189
|
|
|
167
190
|
def write_buffer(buffer)
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
191
|
+
telemetry = @thread_telemetry
|
|
192
|
+
if telemetry.equal?(false)
|
|
193
|
+
telemetry = Thread.current[:pl_telemetry]
|
|
194
|
+
@thread_telemetry = telemetry
|
|
195
|
+
end
|
|
196
|
+
if telemetry
|
|
197
|
+
ticket = telemetry.start(:write_time)
|
|
198
|
+
@io.write(buffer)
|
|
199
|
+
telemetry.finish(:write_time, ticket)
|
|
200
|
+
else
|
|
201
|
+
@io.write(buffer)
|
|
202
|
+
end
|
|
171
203
|
@bytes_written += buffer.bytesize
|
|
172
204
|
rotate_if_needed
|
|
173
205
|
end
|
|
@@ -192,7 +224,8 @@ module Purplelight
|
|
|
192
224
|
|
|
193
225
|
def next_part_path
|
|
194
226
|
ext = 'jsonl'
|
|
195
|
-
|
|
227
|
+
sequence = @part_sequence ? @part_sequence.next : @file_seq
|
|
228
|
+
filename = format('%<prefix>s-part-%<seq>06d.%<ext>s', prefix: @prefix, seq: sequence, ext: ext)
|
|
196
229
|
filename += '.zst' if @effective_compression.to_s == 'zstd'
|
|
197
230
|
filename += '.gz' if @effective_compression.to_s == 'gzip'
|
|
198
231
|
File.join(@directory, filename)
|