herringbone 0.4.0 → 0.5.0

data/lib/herringbone/writer.rb

@@ -3,10 +3,10 @@
 module Herringbone
   # Writes Parquet files.
   #
-  #   schema = Herringbone::Schema.define do
-  #     int64 :id, null: false
-  #     string :name
-  #     list :tags, :string
+  #   schema = Herringbone::Schema.define do |s|
+  #     s.int64 :id, null: false
+  #     s.string :name
+  #     s.list :tags, :string
   #   end
   #   File.open("out.parquet", "wb") do |file|
   #     Herringbone::Writer.open(file, schema) do |w|
@@ -164,6 +164,7 @@ module Herringbone
       @total_rows = 0
       @pos = 0
       @closed = false
+      @aborted = false
       @bytes_per_row = nil
       @io = check_io!(io)
       write_raw(MAGIC)
@@ -222,21 +223,82 @@ module Herringbone
     # @return [void]
     def flush_row_group
       return if @buffered_rows.zero?
+      write_row_group(@buffered_rows)
+    end
+
+    # A column chunk taken as it is from another file, for #write_row_group
+    #
+    # @!attribute chunk
+    #   @return [Format::ColumnChunk] the chunk's footer entry in the source file
+    # @!attribute start
+    #   @return [Integer] source file offset of the chunk's first page
+    # @!attribute bytes
+    #   @return [String] the chunk's pages, headers included
+    # @!attribute column_index
+    #   @return [String, nil] the chunk's encoded ColumnIndex
+    # @!attribute offset_index
+    #   @return [Format::OffsetIndex, nil] the chunk's OffsetIndex, with source file offsets
+    # @!attribute bloom_filter
+    #   @return [String, nil] the chunk's encoded bloom filter, header included
+    CopiedChunk = Struct.new(:chunk, :start, :bytes, :column_index, :offset_index, :bloom_filter)
+
+    # Internal (used by Redaction): writes a row group of +num_rows+ rows from the buffered values,
+    # except for the columns in +copies+, whose chunks are copied byte for byte from another file
+    # with their offsets rebased. Starts a new row group afterwards.
+    #
+    # @param num_rows [Integer] rows in the row group
+    # @param copies [Hash{Integer => CopiedChunk}] column index => chunk to copy instead of encoding
+    # @param codecs [Hash{Integer => Integer}] column index => codec id for an encoded chunk, instead
+    #   of the +compression:+ option
+    # @param bloom_filters [Hash{Integer => Boolean}] column index => true to give an encoded chunk a
+    #   bloom filter (default settings) when the +bloom_filters:+ option does not ask for one
+    # @param sorting_columns [Array<Format::SortingColumn>, nil] stored on the RowGroup as given
+    # @return [void]
+    # @raise [Error] when the writer is closed
+    def write_row_group(num_rows, copies: {}, codecs: {}, bloom_filters: {}, sorting_columns: nil)
+      raise Error, "Writer is closed" if @closed
       start = @pos
-      chunks = @schema.columns.map { |col| write_column_chunk(col, @buffers[col.index]) }
+      chunks = @schema.columns.map do |col|
+        if (copy = copies[col.index])
+          copy_column_chunk(copy)
+        else
+          bloom = @bloom_filters[col.dotted_path]
+          bloom ||= bloom_filter_settings({}) if bloom_filters[col.index] && BloomFilter::TYPES.include?(col.type)
+          write_column_chunk(col, @buffers[col.index], codec: codecs.fetch(col.index, @codec), bloom: bloom)
+        end
+      end
       @row_groups << Format::RowGroup.new(
         columns: chunks,
         total_byte_size: chunks.sum { |c| c.meta_data.total_uncompressed_size },
-        num_rows: @buffered_rows,
+        num_rows: num_rows,
+        sorting_columns: sorting_columns,
         file_offset: start,
         total_compressed_size: @pos - start,
         ordinal: @row_groups.size
       )
       write_bloom_filters
-      @total_rows += @buffered_rows
+      @total_rows += num_rows
       reset_buffers
     end
 
+    # Internal (used by Redaction): shreds one value per row of a top-level field into the buffers,
+    # for #write_row_group. Other fields are left as they are, so the caller decides which columns
+    # are encoded and which are copied.
+    #
+    # @param name [String] top-level field name
+    # @param values [Array] the field's value for each row
+    # @return [void]
+    # @raise [ArgumentError] when the schema has no such field
+    # @raise [EncodeError] when a value cannot be written to the field
+    def buffer_field(name, values)
+      field = @schema.field(name) or raise ArgumentError, "No such field #{name.inspect}"
+      values.each_with_index do |value, i|
+        shred(field, value, 0, 0)
+      rescue EncodeError => e
+        raise EncodeError.new("Row #{@total_rows + i}: #{e.message}", row: @total_rows + i, column: e.column, value: e.value)
+      end
+    end
+
     # Flushes buffered rows, then writes the page indexes and the footer and flushes the IO.
     # Does nothing when already closed or aborted. The IO itself is not closed.
     #
@@ -267,9 +329,23 @@ module Herringbone
     #
     # @return [void]
     def abort
+      @aborted = true unless @closed
       @closed = true
     end
 
+    # Short summary for the console, without the buffered values
+    #
+    # @return [String] state (open, closed or aborted), rows written including buffered ones,
+    #   row groups flushed and the codec
+    def inspect
+      state = if @aborted then "aborted"
+      elsif @closed then "closed"
+      else "open"
+      end
+      "#<#{self.class.name} #{state} rows_written=#{rows_written} row_groups=#{@row_groups.size} " \
+        "compression=#{Compression::NAMES.fetch(@codec, @codec).inspect}>"
+    end
+
     private
 
     # Pathname responds to #write too (writing a whole file by path), so it is rejected explicitly.
@@ -534,8 +610,11 @@ module Herringbone
     #
     # @param col [Schema::Column] column being written
     # @param buffer [ColumnBuffer] the column's buffered levels and values
+    # @param codec [Integer] codec id to compress the pages with
+    # @param bloom [Hash{Symbol => Numeric, nil}, nil] bloom filter settings (see
+    #   #bloom_filter_settings), nil for no bloom filter
     # @return [Format::ColumnChunk] chunk with its ColumnMetaData, for the row group
-    def write_column_chunk(col, buffer)
+    def write_column_chunk(col, buffer, codec: @codec, bloom: @bloom_filters[col.dotted_path])
       type = col.type
       path = col.dotted_path
       dict_values = nil
@@ -568,7 +647,7 @@ module Herringbone
           type: Format::PageType::DICTIONARY_PAGE,
           dictionary_page_header: Format::DictionaryPageHeader.new(num_values: dict_values.size, encoding: E::PLAIN)
         )
-        uncompressed_total += write_page(header, plain)
+        uncompressed_total += write_page(header, plain, codec: codec)
       end
 
       data_offset = @pos
@@ -602,10 +681,10 @@ module Herringbone
         rep_bytes = max_rep.positive? ? Encodings::RLE.encode_hybrid(reps, max_rep.bit_length) : "".b
         def_bytes = max_def.positive? ? Encodings::RLE.encode_hybrid(defs, max_def.bit_length) : "".b
         uncompressed_total += if @data_page_version == 1
-          write_data_page_v1(n, rep_bytes, def_bytes, encoded, value_encoding)
+          write_data_page_v1(n, rep_bytes, def_bytes, encoded, value_encoding, codec)
         else
           num_rows = max_rep.zero? ? n : reps.count(0)
-          write_data_page_v2(n, n - non_null, num_rows, rep_bytes, def_bytes, encoded, value_encoding)
+          write_data_page_v2(n, n - non_null, num_rows, rep_bytes, def_bytes, encoded, value_encoding, codec)
         end
         pages << PageInfo.new(page_offset, @pos - page_offset, first_row, n - non_null, non_null, range)
         first_row += max_rep.zero? ? n : reps.count(0)
@@ -618,7 +697,7 @@ module Herringbone
         type: type,
         encodings: encodings.uniq,
         path_in_schema: col.path,
-        codec: @codec,
+        codec: codec,
         num_values: buf.defs.size,
         total_uncompressed_size: uncompressed_total,
         total_compressed_size: @pos - chunk_start,
@@ -627,7 +706,7 @@ module Herringbone
         statistics: statistics_for(col, buf.defs, values || indices.uniq.map { |i| dict_values[i] }, order)
       )
       chunk = Format::ColumnChunk.new(file_offset: chunk_start, meta_data: meta)
-      if (bloom = @bloom_filters[path])
+      if bloom
         @pending_bloom_filters << [meta, build_bloom_filter(col, bloom, dict_values, values)]
       end
       @page_indexes << [chunk, column_index_for(col, pages, order), offset_index_for(pages)]
@@ -658,6 +737,37 @@ module Herringbone
       end)
     end
 
+    # Writes a chunk copied from another file. The pages, the ColumnIndex and the bloom filter
+    # hold no file offsets and go out as they are; the ColumnMetaData and the OffsetIndex do, so
+    # those are rebased onto where the chunk lands in this file.
+    #
+    # @param copy [CopiedChunk] the chunk to copy
+    # @return [Format::ColumnChunk] chunk with its ColumnMetaData, for the row group
+    def copy_column_chunk(copy)
+      start = @pos
+      shift = start - copy.start
+      meta = Format::ColumnMetaData.decode(copy.chunk.meta_data.encode).first
+      meta.data_page_offset += shift
+      dict = meta.dictionary_page_offset
+      # Some writers store 0 when there is no dictionary page
+      meta.dictionary_page_offset = dict&.positive? ? dict + shift : nil
+      meta.index_page_offset += shift if meta.index_page_offset
+      meta.total_compressed_size = copy.bytes.bytesize
+      meta.bloom_filter_offset = meta.bloom_filter_length = nil
+      write_raw(copy.bytes)
+      chunk = Format::ColumnChunk.new(file_offset: start, meta_data: meta)
+      @pending_bloom_filters << [meta, copy.bloom_filter] if copy.bloom_filter
+      offset_index = copy.offset_index && Format::OffsetIndex.new(
+        page_locations: copy.offset_index.page_locations.map do |loc|
+          Format::PageLocation.new(offset: loc.offset + shift, compressed_page_size: loc.compressed_page_size,
+            first_row_index: loc.first_row_index)
+        end,
+        unencoded_byte_array_data_bytes: copy.offset_index.unencoded_byte_array_data_bytes
+      )
+      @page_indexes << [chunk, copy.column_index, offset_index]
+      chunk
+    end
+
     # nil when the column has no defined sort order, or a page's values have no min/max (all NaN)
     #
     # @param col [Schema::Column] column the pages belong to
@@ -693,18 +803,20 @@ module Herringbone
       end
     end
 
-    # Page indexes go after the last row group: all column indexes, then all offset indexes
+    # Page indexes go after the last row group: all column indexes, then all offset indexes.
+    # A copied chunk brings its ColumnIndex already encoded, and may come without either index.
     #
     # @return [void]
     def write_page_indexes
       @page_indexes.each do |chunk, column_index, _|
         next unless column_index
-        bytes = column_index.encode
+        bytes = column_index.is_a?(String) ? column_index : column_index.encode
         chunk.column_index_offset = @pos
         chunk.column_index_length = bytes.bytesize
         write_raw(bytes)
       end
       @page_indexes.each do |chunk, _, offset_index|
+        next unless offset_index
         bytes = offset_index.encode
         chunk.offset_index_offset = @pos
         chunk.offset_index_length = bytes.bytesize
@@ -780,12 +892,13 @@ module Herringbone
       filter
     end
 
-    # Bloom filters go right after the row group's column chunks, in column order
+    # Bloom filters go right after the row group's column chunks, in column order. A copied
+    # chunk brings its filter already encoded.
     #
     # @return [void]
     def write_bloom_filters
       @pending_bloom_filters.each do |meta, filter|
-        bytes = filter.encode
+        bytes = filter.is_a?(String) ? filter : filter.encode
         meta.bloom_filter_offset = @pos
         meta.bloom_filter_length = bytes.bytesize
         write_raw(bytes)
@@ -906,9 +1019,10 @@ module Herringbone
     # @param body [String] uncompressed page body
     # @param compressed [String, nil] bytes to write as the page body, when already prepared (v2 data
     #   pages, whose levels stay uncompressed); +body+ is compressed otherwise
+    # @param codec [Integer] codec id to compress +body+ with
     # @return [Integer]
-    def write_page(header, body, compressed = nil)
-      compressed ||= Compression.compress(@codec, body)
+    def write_page(header, body, compressed = nil, codec: @codec)
+      compressed ||= Compression.compress(codec, body)
       header.uncompressed_page_size ||= body.bytesize
       header.compressed_page_size = compressed.bytesize
       header.crc = Zlib.crc32(compressed).then { |c| (c >= 0x8000_0000) ? c - 0x1_0000_0000 : c }
@@ -925,8 +1039,9 @@ module Herringbone
     # @param def_bytes [String] RLE-encoded definition levels, empty when the column has none
     # @param encoded [String] encoded values
     # @param encoding [Integer] encoding id of the values
+    # @param codec [Integer] codec id to compress the page with
     # @return [Integer] uncompressed size including the header
-    def write_data_page_v1(n, rep_bytes, def_bytes, encoded, encoding)
+    def write_data_page_v1(n, rep_bytes, def_bytes, encoded, encoding, codec)
       body = String.new(encoding: Encoding::BINARY)
       body << [rep_bytes.bytesize].pack("V") << rep_bytes unless rep_bytes.empty?
       body << [def_bytes.bytesize].pack("V") << def_bytes unless def_bytes.empty?
@@ -938,7 +1053,7 @@ module Herringbone
           definition_level_encoding: E::RLE, repetition_level_encoding: E::RLE
         )
       )
-      write_page(header, body)
+      write_page(header, body, codec: codec)
     end
 
     # A DATA_PAGE_V2: levels without length prefixes and uncompressed, then the compressed values
@@ -950,9 +1065,10 @@ module Herringbone
     # @param def_bytes [String] RLE-encoded definition levels, empty when the column has none
     # @param encoded [String] encoded values
     # @param encoding [Integer] encoding id of the values
+    # @param codec [Integer] codec id to compress the values with
     # @return [Integer] uncompressed size including the header
-    def write_data_page_v2(n, nulls, rows, rep_bytes, def_bytes, encoded, encoding)
-      compressed = Compression.compress(@codec, encoded)
+    def write_data_page_v2(n, nulls, rows, rep_bytes, def_bytes, encoded, encoding, codec)
+      compressed = Compression.compress(codec, encoded)
       header = Format::PageHeader.new(
         type: Format::PageType::DATA_PAGE_V2,
         uncompressed_page_size: rep_bytes.bytesize + def_bytes.bytesize + encoded.bytesize,
@@ -960,10 +1076,10 @@ module Herringbone
           num_values: n, num_nulls: nulls, num_rows: rows, encoding: encoding,
           definition_levels_byte_length: def_bytes.bytesize,
           repetition_levels_byte_length: rep_bytes.bytesize,
-          is_compressed: @codec != Format::Codec::UNCOMPRESSED
+          is_compressed: codec != Format::Codec::UNCOMPRESSED
         )
       )
-      write_page(header, "".b, rep_bytes + def_bytes + compressed)
+      write_page(header, "".b, rep_bytes + def_bytes + compressed, codec: codec)
     end
 
     # Statistics and column index bounds longer than this are truncated, see #truncate_min

data/lib/herringbone.rb

@@ -64,6 +64,7 @@ require_relative "herringbone/xxhash"
 require_relative "herringbone/bloom_filter"
 require_relative "herringbone/inspector"
 require_relative "herringbone/visualizer"
+require_relative "herringbone/redaction"
 
 module Herringbone
   module_function
@@ -77,7 +78,7 @@ module Herringbone
   # to Writer.
   #
   #   File.open("orders.parquet", "wb") { |f| Herringbone.write(f, Order.where(created_at: 1.year.ago..)) }
-  #   Herringbone.write(io, events.lazy.map(&:to_h)) { json :payload }
+  #   Herringbone.write(io, events.lazy.map(&:to_h)) { |s| s.json :payload }
   #
   # @param io [IO, #write] destination; written sequentially, never closed
   # @param records [Enumerable<Hash, Array, Object>, Class, #find_each] rows (Hashes, Arrays in schema
@@ -97,14 +98,17 @@ module Herringbone
   # @option options [Hash{String => String}] :metadata ({}) key/value metadata for the footer
   # @option options [Boolean, Array<String>, Hash{String => Boolean, Hash}] :bloom_filters (nil)
   #   columns to write split block bloom filters for, see Writer
-  # @yield optional block for Schema.infer (Builder DSL), declaring fields that replace inferred ones;
-  #   ignored when the schema is not inferred
+  # @yield [s] optional, declares fields that replace inferred ones (see Schema.infer); ignored when
+  #   the schema is not inferred
+  # @yieldparam s [Schema::Builder] the builder to declare fields on
+  # @yieldreturn [void]
   # @return [Integer] number of rows written
   # @raise [ArgumentError] when the schema has to be inferred and +records+ is empty or holds Array
-  #   rows, or an option is invalid
+  #   rows, an option is invalid, or the block takes no parameter
   # @raise [EncodeError] when a row does not fit the schema
   # @raise [SchemaMismatch] when a row does not fit the inferred schema; the file is left unfinished
   def write(io, records, schema: nil, **options, &overrides)
+    Schema::Builder.check_block!(overrides, "Herringbone.write(io, rows) { |s| s.json :payload }")
     model = if records.respond_to?(:klass) then records.klass
     elsif records.respond_to?(:columns) && records.respond_to?(:find_each) then records
     end
@@ -112,7 +116,7 @@ module Herringbone
     writer = if schema
       Writer.new(io, schema, **options)
     else
-      fix = "Herringbone.write(io, rows) { %s }"
+      fix = "Herringbone.write(io, rows) { |s| s.%s }"
       InferringWriter.new(io, fix: fix, **options) { |sample| Schema.infer(sample, &overrides) }
     end
     begin
@@ -129,6 +133,44 @@ module Herringbone
     writer.rows_written
   end
 
+  # Rewrites +input+ into +output+ with rows removed or values replaced, building the Redaction from
+  # the block (or taking +redaction+). A shortcut for Redaction#apply.
+  #
+  #   Herringbone.redact(input, output) do |r|
+  #     r.where(user_id: 42).delete
+  #     r.replace(:email) { |email| email && OpenSSL::HMAC.hexdigest("SHA256", KEY, email) }
+  #   end
+  #
+  # @param input [IO, StringIO] the Parquet file, read with #seek and #read; not closed
+  # @param output [IO, #write] destination, written sequentially; not closed
+  # @param redaction [Redaction, nil] the redaction to apply, instead of a block
+  # @param writer_options [Hash{Symbol => Object}] Writer options for re-encoded column chunks, and
+  #   +metadata:+ to replace the footer key/value metadata, see Redaction#apply
+  # @option writer_options [Symbol] :compression (codec of each source chunk) codec for re-encoded chunks
+  # @option writer_options [Boolean, Array<String>, Hash{String => Boolean, Hash}] :bloom_filters (nil)
+  #   columns whose re-encoded chunks get a bloom filter, besides those whose source chunk had one
+  # @option writer_options [Hash{String => String}] :metadata (the input's) footer key/value metadata
+  # @option writer_options [Integer] :page_bytes (1MB) approximate uncompressed data page size
+  # @option writer_options [Integer] :page_rows (20_000) maximum rows per data page
+  # @option writer_options [Integer] :data_page_version (1) 1 or 2
+  # @option writer_options [Boolean, Array<String>] :dictionary (true) see Writer
+  # @option writer_options [Hash{String => Symbol}] :encodings ({}) see Writer
+  # @yield [r] declares the statements on a new Redaction (see Redaction.new)
+  # @yieldparam r [Redaction] the redaction being built
+  # @yieldreturn [void]
+  # @return [Redaction::Report] what was done
+  # @raise [ArgumentError] when given both or neither of +redaction+ and a block, a block that takes
+  #   no parameter, or when the redaction does not fit the file's schema
+  # @raise [EncodeError] when a replacement value cannot be written to its column
+  def redact(input, output, redaction = nil, **writer_options, &block)
+    raise ArgumentError, "Herringbone.redact takes a Redaction or a block, not both" if redaction && block
+    raise ArgumentError, "Herringbone.redact needs a Redaction or a block" unless redaction || block
+    if block&.arity&.zero?
+      raise ArgumentError, "The block receives the redaction: Herringbone.redact(input, output) { |r| r.where(user_id: 42).delete }"
+    end
+    (redaction || Redaction.new(&block)).apply(input, output, **writer_options)
+  end
+
   # Compression codecs this process can read and write, e.g. [:none, :snappy, :gzip, :lz4, :lz4_hadoop, :zstd].
   # :zstd and :brotli are listed when the zstd-ruby / brotli gems are loaded.
   #

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: herringbone
 version: !ruby/object:Gem::Version
-  version: 0.4.0
+  version: 0.5.0
 platform: ruby
 authors:
 - Julik Tarkhanov
@@ -58,6 +58,8 @@ files:
 - lib/herringbone/reader/numo.rb
 - lib/herringbone/reader/page_stream.rb
 - lib/herringbone/reader/scan.rb
+- lib/herringbone/redaction.rb
+- lib/herringbone/redaction/rewriter.rb
 - lib/herringbone/schema.rb
 - lib/herringbone/simple_writer.rb
 - lib/herringbone/thrift.rb