Class: Purplelight::WriterJSONL

Inherits:
Object
  • Object
show all
Defined in:
lib/purplelight/writer_jsonl.rb

Overview

WriterJSONL writes newline-delimited JSON with optional compression.

Defined Under Namespace

Classes: EncodedBatch, PartSequence

Constant Summary collapse

DEFAULT_ROTATE_BYTES = 
256 * 1024 * 1024
DEFAULT_ZSTD_LEVEL = 
3
DEFAULT_GZIP_LEVEL = 
1

Instance Method Summary collapse

Constructor Details

#initialize(directory:, prefix:, compression: :zstd, rotate_bytes: DEFAULT_ROTATE_BYTES, logger: nil, manifest: nil, compression_level: nil, write_chunk_bytes: nil, part_sequence: nil) ⇒ WriterJSONL

Returns a new instance of WriterJSONL.



43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/purplelight/writer_jsonl.rb', line 43

def initialize(directory:, prefix:, compression: :zstd, rotate_bytes: DEFAULT_ROTATE_BYTES, logger: nil,
               manifest: nil, compression_level: nil, write_chunk_bytes: nil, part_sequence: nil)
  @directory = directory
  @prefix = prefix
  @compression = compression
  @rotate_bytes = rotate_bytes
  @logger = logger
  @manifest = manifest
  env_level = ENV['PL_ZSTD_LEVEL']&.to_i
  @compression_level = compression_level || (env_level&.positive? ? env_level : nil)
  @write_chunk_bytes = write_chunk_bytes
  @part_sequence = part_sequence

  @part_index = nil
  @io = nil
  @bytes_written = 0
  @file_seq = manifest ? manifest.parts.length : 0
  @closed = false
  @thread_telemetry = false

  @effective_compression = determine_effective_compression(@compression)
  @json_state = JSON::Ext::Generator::State.new(ascii_only: false, max_nesting: false)
  if @logger
    level_disp = @compression_level
    @logger.info("WriterJSONL using compression='#{@effective_compression}' level='#{level_disp || 'default'}'")
  end
  return unless @effective_compression.to_s != @compression.to_s

  @logger&.warn("requested compression '#{@compression}' not available; using '#{@effective_compression}'")
end

Instance Method Details

#close ⇒ Object 



144
145
146
147
148
149
150
151
152
 
# File 'lib/purplelight/writer_jsonl.rb', line 144

def close
  return if @closed

  if @io
    finalize_current_part!
    @io.close
  end
  @closed = true
end

#rotate_if_needed ⇒ Object 



137
138
139
140
141
142
 
# File 'lib/purplelight/writer_jsonl.rb', line 137

def rotate_if_needed
  return if @rotate_bytes.nil?
  return if @bytes_written < @rotate_bytes

  rotate!
end

#write_many(batch) ⇒ Object 



74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
 
# File 'lib/purplelight/writer_jsonl.rb', line 74

def write_many(batch)
  ensure_open!

  chunk_threshold = @write_chunk_bytes || ENV['PL_WRITE_CHUNK_BYTES']&.to_i || (8 * 1024 * 1024)
  total_bytes = 0
  rows = 0

  if batch.is_a?(EncodedBatch)
    write_buffer(batch.data)
    rows = batch.rows
    total_bytes = batch.bytes
  elsif batch.is_a?(String)
    # Fast path for callers that don't provide row metadata.
    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|
      json = @json_state.generate(doc)
      rows += 1
      bytes = json.bytesize + 1
      chunk << json
      chunk << "\n"
      chunk_bytes += bytes
      next unless chunk_bytes >= chunk_threshold

      write_buffer(chunk)
      total_bytes += chunk_bytes
      chunk = +''
      chunk_bytes = 0
    end
    unless chunk.empty?
      write_buffer(chunk)
      total_bytes += chunk_bytes
    end
  end

  @manifest&.add_progress_to_part!(index: @part_index, rows_delta: rows, bytes_delta: total_bytes)
end