herringbone 0.4.0 → 0.6.0

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,585 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Herringbone
4
+ class Redaction
5
+ # Applies a Redaction to one file. Building it checks the redaction against the file's schema,
6
+ # so mistakes surface before anything is written.
7
+ #
8
+ # Each row group lands in one of three tiers. When no statement can match (by statistics,
9
+ # bloom filters and the page index, or by reading the condition columns), its column chunks
10
+ # are copied byte for byte. When rows match but none is deleted and only leaf columns change,
11
+ # the changed chunks are re-encoded and the others copied. Otherwise the whole row group is
12
+ # rewritten, keeping its boundary: one row group in, one (or none) out.
13
+ class Rewriter
14
+ # A replaced column, resolved against the file's schema
15
+ #
16
+ # @!attribute name
17
+ # @return [String] the column as named in the statement
18
+ # @!attribute path
19
+ # @return [Array<String>] top-level field name, then struct member names
20
+ # @!attribute field
21
+ # @return [Schema::Field] the replaced field
22
+ # @!attribute column
23
+ # @return [Schema::Column, nil] the leaf column, nil when the field is nested
24
+ Target = Struct.new(:name, :path, :field, :column)
25
+
26
+ # Writer options that make no sense here, since row groups keep their boundaries
27
+ ROW_GROUP_OPTIONS = %i[row_group_bytes row_group_rows].freeze
28
+ # Footer metadata that describes the columns, and goes stale when some are dropped
29
+ COLUMN_METADATA_KEYS = %w[ARROW:schema pandas].freeze
30
+ # Bytes read at a time when a page header runs past what was read of a chunk
31
+ READ_MORE = 64 * 1024
32
+ # What #current returns for a struct member of a null struct
33
+ ABSENT = Object.new.freeze
34
+
35
+ # @param redaction [Redaction] the statements and drops to apply
36
+ # @param io [IO, StringIO] the input file, read with #seek and #read
37
+ # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted input, see Reader.new
38
+ # @raise [ArgumentError] when a statement or drop does not fit the file's schema
39
+ # @raise [FormatError] when the footer cannot be read
40
+ # @raise [DecryptionError] when the footer is encrypted and cannot be decrypted
41
+ def initialize(redaction, io, decryption: nil)
42
+ @io = io
43
+ @reader = Reader.new(io, decryption: decryption)
44
+ @schema = @reader.schema
45
+ @statements = redaction.statements
46
+ @filters = @statements.map { |s| s.where && Reader::Filter.new(@schema, s.where) }
47
+ @targets = @statements.map { |s| s.targets.map { |name| target(name, s) } }
48
+ @row_blocks = @statements.map { |s| s.block && wants_row?(s.block) }
49
+ @drops = redaction.drops.uniq.map { |name| resolve(name, "drop")[1] }
50
+ check_drops!
51
+ @output_schema = output_schema
52
+ @output_columns = @output_schema.columns.map { |c| @schema.column(c.path) }
53
+ @first_rows = @reader.row_groups.each_with_object([0]) { |rg, firsts| firsts << firsts.last + rg.num_rows }
54
+ @rows_read = 0
55
+ @data = {}
56
+ end
57
+
58
+ # @return [Boolean] see Redaction#affects?
59
+ def affects?
60
+ return true unless @drops.empty?
61
+ @reader.row_groups.each_index.any? do |i|
62
+ @data = {}
63
+ candidates = candidates(i)
64
+ next false if candidates.empty?
65
+ candidates.any? { |k| @filters[k].nil? } || any_match?(i, candidates)
66
+ end
67
+ end
68
+
69
+ # @param output [IO, #write] destination
70
+ # @param options [Hash{Symbol => Object}] Writer options for re-encoded chunks, and +metadata:+
71
+ # @option options [Symbol] :compression (codec of each source chunk) codec for re-encoded chunks
72
+ # @option options [Integer, nil] :compression_level (nil) level for that codec, see Writer
73
+ # @option options [Boolean, Array<String>, Hash{String => Boolean, Hash}] :bloom_filters (nil)
74
+ # columns whose re-encoded chunks get a bloom filter, besides those whose source chunk had one
75
+ # @option options [Hash{String => String}] :metadata (the input's) footer key/value metadata
76
+ # @option options [Integer] :page_bytes (1MB) approximate uncompressed data page size
77
+ # @option options [Integer] :page_rows (20_000) maximum rows per data page
78
+ # @option options [Integer] :data_page_version (1) 1 or 2
79
+ # @option options [Boolean, Array<String>] :dictionary (true) see Writer
80
+ # @option options [Hash{String => Symbol}] :encodings ({}) see Writer
81
+ # @option options [EncryptionConfiguration, Hash{Symbol => Object}, false] :encryption (as the input) see Writer
82
+ # @return [Report]
83
+ # @raise [ArgumentError] for +row_group_bytes:+ / +row_group_rows:+ or an invalid writer option
84
+ # @raise [EncodeError] when a replacement value cannot be written
85
+ # @raise [DecryptionError] when the output is to be encrypted like the input but a key is missing
86
+ def apply(output, **options)
87
+ bad = options.keys & ROW_GROUP_OPTIONS
88
+ raise ArgumentError, "#{bad.join(", ")}: a redaction keeps the row groups of the input" unless bad.empty?
89
+ @keep_codecs = !options.key?(:compression)
90
+ options = {metadata: copied_metadata}.merge(options)
91
+ options[:encryption] = input_encryption unless options.key?(:encryption)
92
+ writer = Writer.new(output, @output_schema, row_group_bytes: 1 << 62, **options)
93
+ @report = Report.new(rows_read: 0, rows_deleted: 0, rows_changed: 0, row_groups: {copied: 0, rewritten: 0})
94
+ begin
95
+ @reader.row_groups.each_index { |i| redact_row_group(i, writer) }
96
+ rescue Exception # rubocop:disable Lint/RescueException -- also abort on Interrupt
97
+ writer.abort
98
+ raise
99
+ end
100
+ writer.close
101
+ @report.rows_read = @rows_read
102
+ @report
103
+ end
104
+
105
+ private
106
+
107
+ # @param i [Integer] row group index
108
+ # @param writer [Writer] the output
109
+ # @return [void]
110
+ def redact_row_group(i, writer)
111
+ @data = {}
112
+ candidates = candidates(i)
113
+ if candidates.empty? || (candidates.all? { |k| @filters[k] } && !any_match?(i, candidates))
114
+ return copy_row_group(i, writer)
115
+ end
116
+ deleted, changed = run_statements(i)
117
+ if deleted.none? && changed.empty?
118
+ copy_row_group(i, writer)
119
+ elsif deleted.none? && changed.all?(&:column)
120
+ rewrite_columns(i, writer, changed)
121
+ else
122
+ rewrite_row_group(i, writer, deleted, changed)
123
+ end
124
+ end
125
+
126
+ # Statements that may touch row group +i+: those without conditions, and those whose
127
+ # conditions the statistics, bloom filters and page index do not rule out
128
+ #
129
+ # @param i [Integer] row group index
130
+ # @return [Array<Integer>] statement indexes
131
+ def candidates(i)
132
+ return [] if rows(i).zero?
133
+ @statements.each_index.select do |k|
134
+ filter = @filters[k]
135
+ filter.nil? || (filter.row_group_may_match?(@reader, i) && !filter.page_ranges(@reader, i).empty?)
136
+ end
137
+ end
138
+
139
+ # Reads the condition columns of row group +i+ and checks them against the original values.
140
+ # When no row matches any statement, no statement changes anything (later statements see
141
+ # what earlier ones changed, and none did), so the row group can be copied.
142
+ #
143
+ # @param i [Integer] row group index
144
+ # @param candidates [Array<Integer>] indexes of statements, all with conditions
145
+ # @return [Boolean] whether some row matches some statement
146
+ def any_match?(i, candidates)
147
+ fields = candidates.flat_map { |k| @filters[k].fields }.uniq
148
+ load(i, fields.map(&:name))
149
+ data = fields.map { |f| @data[f.name] }
150
+ candidates.any? { |k| !@filters[k].matching_rows(data, fields, rows(i), false).empty? }
151
+ end
152
+
153
+ # Applies the statements to every row of row group +i+, in place in @data
154
+ #
155
+ # @param i [Integer] row group index
156
+ # @return [Array(Array<Boolean>, Array<Target>)] deleted flag per row, and the targets whose
157
+ # values changed
158
+ def run_statements(i)
159
+ names = if @row_blocks.any?
160
+ @schema.fields.map(&:name)
161
+ else
162
+ @filters.compact.flat_map { |f| f.fields.map(&:name) } + @targets.flatten.map { |t| t.path.first }
163
+ end
164
+ load(i, names.uniq)
165
+ n = rows(i)
166
+ deleted = Array.new(n, false)
167
+ changed = {}
168
+ n.times do |r|
169
+ row_changed = false
170
+ @statements.each_with_index do |statement, k|
171
+ filter = @filters[k]
172
+ next if filter && !row_matches?(filter, r)
173
+ if statement.kind == :delete
174
+ deleted[r] = true
175
+ break
176
+ end
177
+ @targets[k].each do |t|
178
+ old = current(t, r)
179
+ next if old.equal?(ABSENT)
180
+ value = if @row_blocks[k] then statement.block.call(old, row_hash(r))
181
+ elsif statement.block then statement.block.call(old)
182
+ else statement.constants[t.name]
183
+ end
184
+ next unless value != old
185
+ assign(t, r, value)
186
+ changed[t.path] ||= t
187
+ row_changed = true
188
+ end
189
+ end
190
+ if deleted[r]
191
+ @report.rows_deleted += 1
192
+ elsif row_changed
193
+ @report.rows_changed += 1
194
+ end
195
+ end
196
+ [deleted, changed.values]
197
+ end
198
+
199
+ # @param filter [Reader::Filter] conditions of a statement
200
+ # @param r [Integer] row index within the row group
201
+ # @return [Boolean] whether the row, as earlier statements left it, matches all conditions
202
+ def row_matches?(filter, r)
203
+ filter.conditions.all? do |c|
204
+ value = c.path.reduce(@data[c.field.name][r]) { |h, key| h.is_a?(Hash) ? h[key] : nil }
205
+ Reader::Filter.matches?(c.test, value)
206
+ end
207
+ end
208
+
209
+ # @param t [Target] replaced column
210
+ # @param r [Integer] row index within the row group
211
+ # @return [Object] the column's current value, or ABSENT when a struct holding it is null
212
+ def current(t, r)
213
+ value = @data[t.path.first][r]
214
+ t.path.drop(1).each do |key|
215
+ return ABSENT unless value.is_a?(Hash)
216
+ value = value[key]
217
+ end
218
+ value
219
+ end
220
+
221
+ # Stores a value, copying the struct Hashes on the way so nothing else shares them
222
+ #
223
+ # @param t [Target] replaced column
224
+ # @param r [Integer] row index within the row group
225
+ # @param value [Object] the new value
226
+ # @return [void]
227
+ def assign(t, r, value)
228
+ column = @data[t.path.first]
229
+ return column[r] = value if t.path.size == 1
230
+ hash = column[r] = column[r].dup
231
+ t.path[1...-1].each { |key| hash = hash[key] = hash[key].dup }
232
+ hash[t.path.last] = value
233
+ end
234
+
235
+ # @param r [Integer] row index within the row group
236
+ # @return [Hash{String => Object}] the whole row as it stands, keyed by top-level field name
237
+ def row_hash(r)
238
+ @schema.fields.to_h { |f| [f.name, @data[f.name][r]] }
239
+ end
240
+
241
+ # Tier 1: every kept column chunk is copied as it is, except those encrypted in the input or
242
+ # the output, which are encoded again
243
+ #
244
+ # @param i [Integer] row group index
245
+ # @param writer [Writer] the output
246
+ # @return [void]
247
+ def copy_row_group(i, writer)
248
+ unless encrypted_columns(i, writer).empty?
249
+ rewrite_columns(i, writer, [])
250
+ @report.row_groups[:copied] += 1
251
+ return
252
+ end
253
+ copies = @output_columns.each_with_index.to_h { |col, j| [j, copied_chunk(i, col)] }
254
+ writer.write_row_group(rows(i), copies: copies, sorting_columns: sorting_columns(i, []))
255
+ @report.row_groups[:copied] += 1
256
+ end
257
+
258
+ # Tier 2: the chunks of the changed leaf columns (and of encrypted ones) are encoded again,
259
+ # the rest are copied
260
+ #
261
+ # @param i [Integer] row group index
262
+ # @param writer [Writer] the output
263
+ # @param changed [Array<Target>] leaf targets whose values changed
264
+ # @return [void]
265
+ def rewrite_columns(i, writer, changed)
266
+ rewritten = changed.map { |t| t.column.index }
267
+ encoded = rewritten | encrypted_columns(i, writer)
268
+ names = encoded.map { |index| @schema.columns[index].path.first }.uniq
269
+ load(i, names)
270
+ names.each { |name| writer.buffer_field(name, @data[name]) }
271
+ copies = {}
272
+ codecs = {}
273
+ blooms = {}
274
+ @output_columns.each_with_index do |col, j|
275
+ if encoded.include?(col.index)
276
+ chunk_settings(i, col, j, codecs, blooms)
277
+ else
278
+ copies[j] = copied_chunk(i, col)
279
+ end
280
+ end
281
+ writer.write_row_group(rows(i), copies: copies, codecs: codecs, bloom_filters: blooms,
282
+ sorting_columns: sorting_columns(i, rewritten))
283
+ @report.row_groups[:rewritten] += 1 unless changed.empty?
284
+ end
285
+
286
+ # Kept columns whose chunk in row group +i+ cannot be copied, because it is encrypted in the
287
+ # input (its AAD names the file and the chunk's place in it) or is to be encrypted
288
+ #
289
+ # @param i [Integer] row group index
290
+ # @param writer [Writer] the output
291
+ # @return [Array<Integer>] input column indexes
292
+ def encrypted_columns(i, writer)
293
+ chunks = @reader.row_groups[i].columns
294
+ @output_columns.each_with_index.filter_map do |col, j|
295
+ col.index if chunks.fetch(col.index).crypto_metadata || writer.encrypted_column?(j)
296
+ end
297
+ end
298
+
299
+ # The Writer +encryption:+ option that encrypts the output like the input: the same
300
+ # algorithm, footer mode, AAD prefix and keys, for the columns that are kept
301
+ #
302
+ # @return [EncryptionConfiguration, nil] nil for a plaintext input
303
+ # @raise [DecryptionError] when a key of the input was not given
304
+ def input_encryption
305
+ decryptor = @reader.decryptor or return nil
306
+ chunks = @reader.row_groups.first&.columns || []
307
+ columns = @output_columns.each_with_index.filter_map do |col, j|
308
+ chunk = chunks[col.index]
309
+ crypto = chunk&.crypto_metadata or next
310
+ path = @output_schema.columns[j].dotted_path
311
+ with_column_key = crypto.encryption_with_column_key
312
+ next [path, :footer] unless with_column_key
313
+ key = decryptor.chunk_key(chunk, col.dotted_path)
314
+ key or raise DecryptionError, "The output is encrypted like the input, which needs the key of #{col.dotted_path}: " \
315
+ "pass it in decryption:, or pass encryption: for the output"
316
+ [path, {key: key, key_metadata: with_column_key.key_metadata}]
317
+ end
318
+ uniform = !columns.empty? && columns.size == @output_columns.size && columns.all? { |_, v| v == :footer }
319
+ decryptor.writer_settings(uniform ? nil : columns.to_h)
320
+ end
321
+
322
+ # Tier 3: the remaining rows are encoded again, every column of them
323
+ #
324
+ # @param i [Integer] row group index
325
+ # @param writer [Writer] the output
326
+ # @param deleted [Array<Boolean>] deleted flag per row
327
+ # @param changed [Array<Target>] targets whose values changed
328
+ # @return [void]
329
+ def rewrite_row_group(i, writer, deleted, changed)
330
+ @report.row_groups[:rewritten] += 1
331
+ kept = deleted.count(false)
332
+ return if kept.zero?
333
+ load(i, @output_schema.fields.map(&:name))
334
+ @output_schema.fields.each do |f|
335
+ values = @data[f.name]
336
+ values = values.reject.with_index { |_, r| deleted[r] } if kept < values.size
337
+ writer.buffer_field(f.name, values)
338
+ end
339
+ codecs = {}
340
+ blooms = {}
341
+ @output_columns.each_with_index { |col, j| chunk_settings(i, col, j, codecs, blooms) }
342
+ rewritten = changed.flat_map { |t| t.field.leaves.map(&:index) }
343
+ writer.write_row_group(kept, codecs: codecs, bloom_filters: blooms, sorting_columns: sorting_columns(i, rewritten))
344
+ end
345
+
346
+ # Codec and bloom filter of a re-encoded chunk follow the source chunk: the same codec
347
+ # (unless +compression:+ was given; LZO cannot be written and becomes Snappy), and a bloom
348
+ # filter when the source chunk had one
349
+ #
350
+ # @param i [Integer] row group index
351
+ # @param col [Schema::Column] the input column
352
+ # @param j [Integer] the output column index
353
+ # @param codecs [Hash{Integer => Integer}] collects output column index => codec id
354
+ # @param blooms [Hash{Integer => Boolean}] collects output column index => true
355
+ # @return [void]
356
+ def chunk_settings(i, col, j, codecs, blooms)
357
+ meta = @reader.row_groups[i].columns.fetch(col.index).meta_data
358
+ return unless meta
359
+ codecs[j] = (meta.codec == Format::Codec::LZO) ? Format::Codec::SNAPPY : meta.codec if @keep_codecs
360
+ blooms[j] = true if meta.bloom_filter_offset
361
+ end
362
+
363
+ # The row group's sorting columns that still hold: dropped or changed columns end the list,
364
+ # since the columns after them were only sorted within their runs
365
+ #
366
+ # @param i [Integer] row group index
367
+ # @param rewritten [Array<Integer>] indexes of input columns whose values changed
368
+ # @return [Array<Format::SortingColumn>, nil]
369
+ def sorting_columns(i, rewritten)
370
+ out = []
371
+ (@reader.row_groups[i].sorting_columns || []).each do |sc|
372
+ col = @schema.columns[sc.column_idx]
373
+ j = col && !rewritten.include?(col.index) && @output_columns.index(col)
374
+ break unless j
375
+ out << Format::SortingColumn.new(column_idx: j, descending: sc.descending, nulls_first: sc.nulls_first)
376
+ end
377
+ out.empty? ? nil : out
378
+ end
379
+
380
+ # @param i [Integer] row group index
381
+ # @param col [Schema::Column] the input column
382
+ # @return [Writer::CopiedChunk] the chunk's pages, page index and bloom filter
383
+ # @raise [UnsupportedError] for an encrypted chunk or one stored in another file
384
+ # @raise [FormatError] when a page header is corrupt or the chunk overruns the file
385
+ def copied_chunk(i, col)
386
+ chunk = @reader.row_groups[i].columns.fetch(col.index)
387
+ meta = chunk.meta_data or raise UnsupportedError, "Column chunk without metadata (encrypted?)"
388
+ raise UnsupportedError, "Column chunks in external files are not supported" if chunk.file_path
389
+ column_index, offset_index = @reader.page_index(i, col)
390
+ start, bytes = chunk_bytes(col, meta, offset_index)
391
+ column_index &&= read_at(chunk.column_index_offset, chunk.column_index_length)
392
+ Writer::CopiedChunk.new(chunk, start, bytes, column_index, offset_index, bloom_bytes(i, col, meta))
393
+ end
394
+
395
+ # The chunk's pages. They are walked header by header rather than trusting
396
+ # total_compressed_size, which some writers get wrong: copying too little breaks the
397
+ # chunk, and copying too much could carry over bytes of a neighbouring chunk.
398
+ #
399
+ # @param col [Schema::Column] the input column
400
+ # @param meta [Format::ColumnMetaData] the chunk's metadata
401
+ # @param offset_index [Format::OffsetIndex, nil] the chunk's OffsetIndex, whose pages are
402
+ # included even when num_values is reached before them
403
+ # @return [Array(Integer, String)] file offset of the first page, and the pages' bytes
404
+ # @raise [FormatError] when a page header is corrupt or the chunk overruns the file
405
+ def chunk_bytes(col, meta, offset_index)
406
+ start = meta.data_page_offset
407
+ dict = meta.dictionary_page_offset
408
+ start = dict if dict&.positive? && dict < start
409
+ last = offset_index&.page_locations&.last
410
+ min_end = last ? last.offset + last.compressed_page_size - start : 0
411
+ @io.seek(start)
412
+ buf = (@io.read(meta.total_compressed_size) || "").b
413
+ pos = 0
414
+ seen = 0
415
+ while seen < meta.num_values || pos < min_end
416
+ begin
417
+ header, body = Format::PageHeader.decode(buf, pos)
418
+ rescue Thrift::Error
419
+ raise FormatError, "Corrupt page header in #{col.dotted_path}" unless read_more(buf, start, READ_MORE)
420
+ retry
421
+ end
422
+ size = header.compressed_page_size
423
+ raise FormatError, "Negative page size in #{col.dotted_path}" if size.nil? || size.negative?
424
+ short = body + size - buf.bytesize
425
+ if short.positive? && read_more(buf, start, short) < short
426
+ raise FormatError, "Column #{col.dotted_path}: page overruns the file"
427
+ end
428
+ pos = body + size
429
+ seen += header.data_page_header&.num_values || header.data_page_header_v2&.num_values || 0
430
+ end
431
+ [start, (pos == buf.bytesize) ? buf : buf.byteslice(0, pos)]
432
+ end
433
+
434
+ # Appends up to +count+ bytes that follow +buf+ in the file
435
+ #
436
+ # @param buf [String] bytes read from +start+ on; appended to
437
+ # @param start [Integer] file offset of +buf+
438
+ # @param count [Integer] bytes wanted
439
+ # @return [Integer] bytes appended, 0 at the end of the file
440
+ def read_more(buf, start, count)
441
+ @io.seek(start + buf.bytesize)
442
+ more = @io.read(count)
443
+ return 0 if more.nil?
444
+ buf << more.b
445
+ more.bytesize
446
+ end
447
+
448
+ # @param offset [Integer, nil] file offset
449
+ # @param length [Integer, nil] byte count
450
+ # @return [String, nil] the bytes, nil when absent or cut off
451
+ def read_at(offset, length)
452
+ return nil unless offset && length&.positive?
453
+ @io.seek(offset)
454
+ bytes = @io.read(length)
455
+ (bytes&.bytesize == length) ? bytes.b : nil
456
+ end
457
+
458
+ # The chunk's bloom filter as stored. Filters written without a length (older writers) are
459
+ # decoded and encoded again. A filter that cannot be read is left out, which only costs
460
+ # pruning.
461
+ #
462
+ # @param i [Integer] row group index
463
+ # @param col [Schema::Column] the input column
464
+ # @param meta [Format::ColumnMetaData] the chunk's metadata
465
+ # @return [String, nil] header and bitset
466
+ def bloom_bytes(i, col, meta)
467
+ return nil unless meta.bloom_filter_offset
468
+ return read_at(meta.bloom_filter_offset, meta.bloom_filter_length) if meta.bloom_filter_length
469
+ @reader.bloom_filter(i, col)&.encode
470
+ rescue FormatError
471
+ nil
472
+ end
473
+
474
+ # Reads top-level fields of row group +i+ into @data, unless they are there already
475
+ #
476
+ # @param i [Integer] row group index
477
+ # @param names [Array<String>] top-level field names
478
+ # @return [void]
479
+ def load(i, names)
480
+ missing = names - @data.keys
481
+ return if missing.empty?
482
+ n = rows(i)
483
+ @rows_read += n if @data.empty?
484
+ @data.merge!(@reader.read(as: :columns, columns: missing, from: @first_rows[i], limit: n))
485
+ end
486
+
487
+ # @param i [Integer] row group index
488
+ # @return [Integer] rows in the row group
489
+ def rows(i) = @reader.row_groups[i].num_rows
490
+
491
+ # @return [Hash{String => String, nil}] the input's key/value metadata, without the keys that
492
+ # describe the columns when some are dropped
493
+ def copied_metadata
494
+ meta = @reader.metadata
495
+ @drops.empty? ? meta : meta.except(*COLUMN_METADATA_KEYS)
496
+ end
497
+
498
+ # @param block [Proc] a replace block
499
+ # @return [Boolean] whether the block takes the row as a second parameter
500
+ def wants_row?(block)
501
+ params = block.parameters
502
+ params.any? { |type, _| type == :rest } || params.count { |type, _| type == :req || type == :opt } >= 2
503
+ end
504
+
505
+ # Resolves a column named in a statement: a top-level field, or a struct member by dotted path
506
+ #
507
+ # @param name [String] the column as named
508
+ # @param verb [String] "replace" or "drop", for error messages
509
+ # @return [Array(Schema::Field, Array<String>)] the field and its path of names
510
+ # @raise [ArgumentError] when there is no such column, or it is inside a list or map
511
+ def resolve(name, verb)
512
+ field = @schema.field(name)
513
+ return [field, [name]] if field
514
+ parts = name.split(".")
515
+ field = @schema.field(parts.first) or raise ArgumentError, "#{verb}: no such column #{name.inspect}"
516
+ parts.drop(1).each_with_index do |part, depth|
517
+ case field.kind
518
+ when :struct
519
+ field = field.children_by_name[part] or raise ArgumentError, "#{verb}: no such column #{name.inspect}"
520
+ when :list, :map
521
+ whole = parts.first(depth + 1).join(".")
522
+ container = (field.kind == :map) ? "Hash" : "Array"
523
+ hint = (verb == "replace") ? " with a block that receives its #{container}" : ""
524
+ raise ArgumentError, "#{verb}: #{name} is inside the #{field.kind} #{whole}; #{verb} #{whole} instead#{hint}"
525
+ else
526
+ raise ArgumentError, "#{verb}: no such column #{name.inspect}"
527
+ end
528
+ end
529
+ [field, parts]
530
+ end
531
+
532
+ # @param name [String] the column as named in the statement
533
+ # @param statement [Statement] the replace statement
534
+ # @return [Target]
535
+ # @raise [ArgumentError] when the column does not exist, is inside a list or map, or a
536
+ # constant cannot be stored in it (nil in a required column, a value of the wrong type)
537
+ def target(name, statement)
538
+ field, path = resolve(name, "replace")
539
+ if statement.block.nil?
540
+ value = statement.constants[name]
541
+ if value.nil?
542
+ raise ArgumentError, "replace: #{name} is required (null: false) and cannot be set to nil" unless field.optional
543
+ elsif field.leaf?
544
+ begin
545
+ field.column.encoder.call(value)
546
+ rescue ArgumentError, TypeError, NoMethodError, RangeError => e
547
+ raise ArgumentError, "replace: cannot write #{value.inspect} to #{name}: #{e.message}"
548
+ end
549
+ end
550
+ end
551
+ Target.new(name, path, field, field.leaf? ? field.column : nil)
552
+ end
553
+
554
+ # @return [void]
555
+ # @raise [ArgumentError] when a replaced column is dropped
556
+ def check_drops!
557
+ @targets.flatten.each do |t|
558
+ dropped = @drops.find { |path| t.path.first(path.size) == path }
559
+ raise ArgumentError, "replace: #{t.name} is dropped by drop(#{dropped.join(".").inspect})" if dropped
560
+ end
561
+ end
562
+
563
+ # The input schema without the dropped fields
564
+ #
565
+ # @return [Schema]
566
+ # @raise [ArgumentError] when every column, or every member of a struct, would be dropped
567
+ def output_schema
568
+ return @schema if @drops.empty?
569
+ copy = Schema.from_elements(@schema.to_elements)
570
+ emptied = @drops.filter_map do |path|
571
+ parent = path[0...-1].reduce(copy.root) { |node, name| node&.children&.find { |c| c.name == name } }
572
+ node = parent&.children&.find { |c| c.name == path.last }
573
+ next unless node
574
+ parent.children.delete(node)
575
+ parent if parent.children.empty?
576
+ end
577
+ emptied.each do |node|
578
+ raise ArgumentError, "drop: cannot drop every column" if node.equal?(copy.root)
579
+ raise ArgumentError, "drop: cannot drop every member of #{node.path.join(".")}; drop #{node.path.join(".")} instead"
580
+ end
581
+ Schema.new(copy.root)
582
+ end
583
+ end
584
+ end
585
+ end