herringbone 0.6.1 → 0.8.0

data/lib/herringbone/redaction/rewriter.rb

@@ -10,6 +10,8 @@ module Herringbone
     # are copied byte for byte. When rows match but none is deleted and only leaf columns change,
     # the changed chunks are re-encoded and the others copied. Otherwise the whole row group is
     # rewritten, keeping its boundary: one row group in, one (or none) out.
+    #
+    # @api private
     class Rewriter
       # A replaced column, resolved against the file's schema
       #
@@ -23,24 +25,21 @@ module Herringbone
       #   @return [Schema::Column, nil] the leaf column, nil when the field is nested
       Target = Struct.new(:name, :path, :field, :column)
 
-      # Writer options that make no sense here, since row groups keep their boundaries
-      ROW_GROUP_OPTIONS = %i[row_group_bytes row_group_rows].freeze
-      # Footer metadata that describes the columns, and goes stale when some are dropped
-      COLUMN_METADATA_KEYS = %w[ARROW:schema pandas].freeze
-      # Bytes read at a time when a page header runs past what was read of a chunk
-      READ_MORE = 64 * 1024
       # What #current returns for a struct member of a null struct
       ABSENT = Object.new.freeze
 
       # @param redaction [Redaction] the statements and drops to apply
-      # @param io [IO, StringIO] the input file, read with #seek and #read
-      # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted input, see Reader.new
-      # @raise [ArgumentError] when a statement or drop does not fit the file's schema
+      # @param io_or_reader [IO, StringIO, Reader] the Parquet file, read with #seek and #read, or a
+      #   Reader of it (whose IO and decryption are used)
+      # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted
+      #   file, see Reader.new; not with a Reader
+      # @raise [ArgumentError] when +io_or_reader+ is neither an IO nor a Reader, when a Reader comes
+      #   with +decryption:+, or when a statement or drop does not fit the file's schema
       # @raise [FormatError] when the footer cannot be read
       # @raise [DecryptionError] when the footer is encrypted and cannot be decrypted
-      def initialize(redaction, io, decryption: nil)
-        @io = io
-        @reader = Reader.new(io, decryption: decryption)
+      def initialize(redaction, io_or_reader, decryption: nil)
+        @reader = fixed_options_reader(io_or_reader, decryption)
+        @copier = Reader::ChunkCopier.new(@reader)
         @schema = @reader.schema
         @statements = redaction.statements
         @filters = @statements.map { |s| s.where && Reader::Filter.new(@schema, s.where) }
@@ -66,7 +65,7 @@ module Herringbone
         end
       end
 
-      # @param output [IO, #write] destination
+      # @param output_io [IO, #write] destination
       # @param options [Hash{Symbol => Object}] Writer options for re-encoded chunks, and +metadata:+
       # @option options [Symbol] :compression (codec of each source chunk) codec for re-encoded chunks
       # @option options [Integer, nil] :compression_level (nil) level for that codec, see Writer
@@ -83,27 +82,43 @@ module Herringbone
       # @raise [ArgumentError] for +row_group_bytes:+ / +row_group_rows:+ or an invalid writer option
       # @raise [EncodeError] when a replacement value cannot be written
       # @raise [DecryptionError] when the output is to be encrypted like the input but a key is missing
-      def apply(output, **options)
-        bad = options.keys & ROW_GROUP_OPTIONS
-        raise ArgumentError, "#{bad.join(", ")}: a redaction keeps the row groups of the input" unless bad.empty?
+      def apply(output_io, **options)
         @keep_codecs = !options.key?(:compression)
-        options = {metadata: copied_metadata}.merge(options)
+        options = {metadata: @copier.metadata(@output_schema)}.merge(options)
         options[:encryption] = input_encryption unless options.key?(:encryption)
-        writer = Writer.new(output, @output_schema, row_group_bytes: 1 << 62, **options)
         @report = Report.new(rows_read: 0, rows_deleted: 0, rows_changed: 0, row_groups: {copied: 0, rewritten: 0})
-        begin
+        Writer.open_for_copies(output_io, @output_schema, "a redaction keeps the row groups of the input", **options) do |writer|
           @reader.row_groups.each_index { |i| redact_row_group(i, writer) }
-        rescue Exception # rubocop:disable Lint/RescueException -- also abort on Interrupt
-          writer.abort
-          raise
         end
-        writer.close
         @report.rows_read = @rows_read
         @report
       end
 
       private
 
+      # The caller's Reader may have been built with +keys:+ or +time_zone:+, which change what
+      # #read returns, so values are always read through a Reader with the default options
+      #
+      # @param io_or_reader [IO, StringIO, Reader] the Parquet file, or a Reader of it
+      # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys, for an IO only
+      # @return [Reader]
+      # @raise [ArgumentError] when +io_or_reader+ is neither an IO nor a Reader, or a Reader comes
+      #   with +decryption:+
+      def fixed_options_reader(io_or_reader, decryption)
+        if io_or_reader.is_a?(Reader)
+          if decryption
+            raise ArgumentError, "decryption: cannot be given with a Herringbone::Reader, which " \
+              "already has its decryption (pass decryption: to Reader.new instead)"
+          end
+          return Reader.new(io_or_reader.io, decryption: io_or_reader.decryption)
+        end
+        unless io_or_reader.respond_to?(:seek) && io_or_reader.respond_to?(:read)
+          raise ArgumentError, "io_or_reader must be an IO that supports #seek and #read " \
+            "(e.g. File.open(path, \"rb\")) or a Herringbone::Reader, got #{io_or_reader.class}"
+        end
+        Reader.new(io_or_reader, decryption: decryption)
+      end
+
       # @param i [Integer] row group index
       # @param writer [Writer] the output
       # @return [void]
@@ -250,7 +265,7 @@ module Herringbone
           @report.row_groups[:copied] += 1
           return
         end
-        copies = @output_columns.each_with_index.to_h { |col, j| [j, copied_chunk(i, col)] }
+        copies = @output_columns.each_with_index.to_h { |col, j| [j, @copier.copy(i, col)] }
         writer.write_row_group(rows(i), copies: copies, sorting_columns: sorting_columns(i, []))
         @report.row_groups[:copied] += 1
       end
@@ -275,7 +290,7 @@ module Herringbone
           if encoded.include?(col.index)
             chunk_settings(i, col, j, codecs, blooms)
           else
-            copies[j] = copied_chunk(i, col)
+            copies[j] = @copier.copy(i, col)
           end
         end
         writer.write_row_group(rows(i), copies: copies, codecs: codecs, bloom_filters: blooms,
@@ -303,20 +318,9 @@ module Herringbone
       # @raise [DecryptionError] when a key of the input was not given
       def input_encryption
         decryptor = @reader.decryptor or return nil
-        chunks = @reader.row_groups.first&.columns || []
-        columns = @output_columns.each_with_index.filter_map do |col, j|
-          chunk = chunks[col.index]
-          crypto = chunk&.crypto_metadata or next
-          path = @output_schema.columns[j].dotted_path
-          with_column_key = crypto.encryption_with_column_key
-          next [path, :footer] unless with_column_key
-          key = decryptor.chunk_key(chunk, col.dotted_path)
-          key or raise DecryptionError, "The output is encrypted like the input, which needs the key of #{col.dotted_path}: " \
-            "pass it in decryption:, or pass encryption: for the output"
-          [path, {key: key, key_metadata: with_column_key.key_metadata}]
-        end
-        uniform = !columns.empty? && columns.size == @output_columns.size && columns.all? { |_, v| v == :footer }
-        decryptor.writer_settings(uniform ? nil : columns.to_h)
+        columns = @copier.encrypted_columns(@output_columns.each_with_index.to_h { |col, j| [@output_schema.columns[j].dotted_path, col] })
+        uniform = !columns.empty? && columns.size == @output_columns.size && columns.values.all?(:footer)
+        decryptor.writer_settings(uniform ? nil : columns)
       end
 
       # Tier 3: the remaining rows are encoded again, every column of them
@@ -343,9 +347,8 @@ module Herringbone
         writer.write_row_group(kept, codecs: codecs, bloom_filters: blooms, sorting_columns: sorting_columns(i, rewritten))
       end
 
-      # Codec and bloom filter of a re-encoded chunk follow the source chunk: the same codec
-      # (unless +compression:+ was given; LZO cannot be written and becomes Snappy), and a bloom
-      # filter when the source chunk had one
+      # Codec and bloom filter of a re-encoded chunk follow the source chunk (see
+      # Reader::ChunkCopier#recode_settings), unless +compression:+ was given
       #
       # @param i [Integer] row group index
       # @param col [Schema::Column] the input column
@@ -354,121 +357,19 @@ module Herringbone
       # @param blooms [Hash{Integer => Boolean}] collects output column index => true
       # @return [void]
       def chunk_settings(i, col, j, codecs, blooms)
-        meta = @reader.row_groups[i].columns.fetch(col.index).meta_data
-        return unless meta
-        codecs[j] = (meta.codec == Format::Codec::LZO) ? Format::Codec::SNAPPY : meta.codec if @keep_codecs
-        blooms[j] = true if meta.bloom_filter_offset
+        codec, bloom = @copier.recode_settings(i, col)
+        codecs[j] = codec if codec && @keep_codecs
+        blooms[j] = true if bloom
       end
 
-      # The row group's sorting columns that still hold: dropped or changed columns end the list,
-      # since the columns after them were only sorted within their runs
+      # The row group's sorting columns that still hold: dropped or changed columns end the list
       #
       # @param i [Integer] row group index
       # @param rewritten [Array<Integer>] indexes of input columns whose values changed
       # @return [Array<Format::SortingColumn>, nil]
       def sorting_columns(i, rewritten)
-        out = []
-        (@reader.row_groups[i].sorting_columns || []).each do |sc|
-          col = @schema.columns[sc.column_idx]
-          j = col && !rewritten.include?(col.index) && @output_columns.index(col)
-          break unless j
-          out << Format::SortingColumn.new(column_idx: j, descending: sc.descending, nulls_first: sc.nulls_first)
-        end
-        out.empty? ? nil : out
-      end
-
-      # @param i [Integer] row group index
-      # @param col [Schema::Column] the input column
-      # @return [Writer::CopiedChunk] the chunk's pages, page index and bloom filter
-      # @raise [UnsupportedError] for an encrypted chunk or one stored in another file
-      # @raise [FormatError] when a page header is corrupt or the chunk overruns the file
-      def copied_chunk(i, col)
-        chunk = @reader.row_groups[i].columns.fetch(col.index)
-        meta = chunk.meta_data or raise UnsupportedError, "Column chunk without metadata (encrypted?)"
-        raise UnsupportedError, "Column chunks in external files are not supported" if chunk.file_path
-        column_index, offset_index = @reader.page_index(i, col)
-        start, bytes = chunk_bytes(col, meta, offset_index)
-        column_index &&= read_at(chunk.column_index_offset, chunk.column_index_length)
-        Writer::CopiedChunk.new(chunk, start, bytes, column_index, offset_index, bloom_bytes(i, col, meta))
-      end
-
-      # The chunk's pages. They are walked header by header rather than trusting
-      # total_compressed_size, which some writers get wrong: copying too little breaks the
-      # chunk, and copying too much could carry over bytes of a neighbouring chunk.
-      #
-      # @param col [Schema::Column] the input column
-      # @param meta [Format::ColumnMetaData] the chunk's metadata
-      # @param offset_index [Format::OffsetIndex, nil] the chunk's OffsetIndex, whose pages are
-      #   included even when num_values is reached before them
-      # @return [Array(Integer, String)] file offset of the first page, and the pages' bytes
-      # @raise [FormatError] when a page header is corrupt or the chunk overruns the file
-      def chunk_bytes(col, meta, offset_index)
-        start = meta.data_page_offset
-        dict = meta.dictionary_page_offset
-        start = dict if dict&.positive? && dict < start
-        last = offset_index&.page_locations&.last
-        min_end = last ? last.offset + last.compressed_page_size - start : 0
-        @io.seek(start)
-        buf = (@io.read(meta.total_compressed_size) || "").b
-        pos = 0
-        seen = 0
-        while seen < meta.num_values || pos < min_end
-          begin
-            header, body = Format::PageHeader.decode(buf, pos)
-          rescue Thrift::Error
-            raise FormatError, "Corrupt page header in #{col.dotted_path}" unless read_more(buf, start, READ_MORE)
-            retry
-          end
-          size = header.compressed_page_size
-          raise FormatError, "Negative page size in #{col.dotted_path}" if size.nil? || size.negative?
-          short = body + size - buf.bytesize
-          if short.positive? && read_more(buf, start, short) < short
-            raise FormatError, "Column #{col.dotted_path}: page overruns the file"
-          end
-          pos = body + size
-          seen += header.data_page_header&.num_values || header.data_page_header_v2&.num_values || 0
-        end
-        [start, (pos == buf.bytesize) ? buf : buf.byteslice(0, pos)]
-      end
-
-      # Appends up to +count+ bytes that follow +buf+ in the file
-      #
-      # @param buf [String] bytes read from +start+ on; appended to
-      # @param start [Integer] file offset of +buf+
-      # @param count [Integer] bytes wanted
-      # @return [Integer] bytes appended, 0 at the end of the file
-      def read_more(buf, start, count)
-        @io.seek(start + buf.bytesize)
-        more = @io.read(count)
-        return 0 if more.nil?
-        buf << more.b
-        more.bytesize
-      end
-
-      # @param offset [Integer, nil] file offset
-      # @param length [Integer, nil] byte count
-      # @return [String, nil] the bytes, nil when absent or cut off
-      def read_at(offset, length)
-        return nil unless offset && length&.positive?
-        @io.seek(offset)
-        bytes = @io.read(length)
-        (bytes&.bytesize == length) ? bytes.b : nil
-      end
-
-      # The chunk's bloom filter as stored. Filters written without a length (older writers) are
-      # decoded and encoded again. A filter that cannot be read is left out, which only costs
-      # pruning.
-      #
-      # @param i [Integer] row group index
-      # @param col [Schema::Column] the input column
-      # @param meta [Format::ColumnMetaData] the chunk's metadata
-      # @return [String, nil] header and bitset
-      def bloom_bytes(i, col, meta)
-        return nil unless meta.bloom_filter_offset
-        return read_at(meta.bloom_filter_offset, meta.bloom_filter_length) if meta.bloom_filter_length
-        @reader.bloom_filter(i, col)&.encode
-      rescue FormatError
-        nil
+        kept = @output_columns.each_with_index.to_h { |col, j| [col.index, j] }
+        @copier.sorting_columns(i, kept.except(*rewritten))
       end
 
       # Reads top-level fields of row group +i+ into @data, unless they are there already
@@ -488,13 +389,6 @@ module Herringbone
       # @return [Integer] rows in the row group
       def rows(i) = @reader.row_groups[i].num_rows
 
-      # @return [Hash{String => String, nil}] the input's key/value metadata, without the keys that
-      #   describe the columns when some are dropped
-      def copied_metadata
-        meta = @reader.metadata
-        @drops.empty? ? meta : meta.except(*COLUMN_METADATA_KEYS)
-      end
-
       # @param block [Proc] a replace block
       # @return [Boolean] whether the block takes the row as a second parameter
       def wants_row?(block)

data/lib/herringbone/redaction.rb

@@ -11,8 +11,8 @@ module Herringbone
   #     r.replace(:phone) { |phone| phone&.gsub(/\d(?=\d{3})/, "#") }
   #     r.drop :ssn
   #   end
-  #   File.open("in.parquet", "rb") do |input|
-  #     File.open("out.parquet", "wb") { |output| forget.apply(input, output) }
+  #   File.open("in.parquet", "rb") do |io|
+  #     File.open("out.parquet", "wb") { |output_io| forget.apply(io, output_io) }
   #   end
   #
   # Statements apply in declared order, per row: a deleted row is gone for later statements, and
@@ -187,39 +187,48 @@ module Herringbone
     #
     # @param statement [Statement] the statement to add
     # @return [Redaction] self
+    # @api private
     def add_statement(statement)
       @statements << statement
       self
     end
 
-    # Whether #apply would change +input_io+: true when a column is dropped, when a replace without
+    # Whether #apply would change +io_or_reader+: true when a column is dropped, when a replace without
     # #where meets a non-empty file, or when some row matches a #where. Reads as little as it can:
     # row groups are ruled out with statistics, bloom filters and the page index first, and only
     # the columns the conditions name are read for the rest.
     #
-    # @param input_io [IO, StringIO] the Parquet file, read with #seek and #read
-    # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted file, see Reader.new
+    # @param io_or_reader [IO, StringIO, Reader] the Parquet file, read with #seek and #read, or a
+    #   Reader of it, whose IO and decryption are used
+    # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted
+    #   file, see Reader.new; not with a Reader, which has its own
     # @return [Boolean]
-    # @raise [ArgumentError] when the redaction does not fit the file's schema
+    # @raise [ArgumentError] when the redaction does not fit the file's schema, or a Reader comes
+    #   with +decryption:+
     # @raise [DecryptionError] when the file is encrypted and a key it needs was not given
-    def affects?(input_io, decryption: nil)
+    def affects?(io_or_reader, decryption: nil)
       check_scopes!
-      Rewriter.new(self, input_io, decryption: decryption).affects?
+      Rewriter.new(self, io_or_reader, decryption: decryption).affects?
     end
 
-    # Writes the redacted copy of +input_io+ to +output_io+. The redaction is checked against the
+    # Writes the redacted copy of +io_or_reader+ to +output_io+. The redaction is checked against the
     # file's schema first, so a missing column, a nil for a required column or a #where without a
     # verb raise before anything is written. The output is left unfinished (no footer) when an
     # error happens later, for instance a block returning a value its column cannot store.
     #
+    # A Reader is read with its own IO and decryption, but always with the default +keys:+ and
+    # +time_zone:+, whatever it was built with.
+    #
     # An encrypted input (opened with +decryption:+) is written encrypted the same way: same
     # algorithm, footer mode, AAD prefix, keys and key metadata. +encryption:+ (see Writer) writes
     # it with other settings, and +encryption: false+ writes a plaintext file. Column chunks that
     # are encrypted in either file are re-encoded rather than copied.
     #
-    # @param input_io [IO, StringIO] the Parquet file, read with #seek and #read; not closed
+    # @param io_or_reader [IO, StringIO, Reader] the Parquet file, read with #seek and #read, or a
+    #   Reader of it; not closed
     # @param output_io [IO, #write] destination, written sequentially; not closed
-    # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted input, see Reader.new
+    # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, nil] keys of an encrypted
+    #   input, see Reader.new; not with a Reader, which has its own
     # @param writer_options [Hash{Symbol => Object}] Writer options for the re-encoded column chunks
     #   (+compression:+, +bloom_filters:+, +page_rows:+, +dictionary:+...), and +metadata:+ to
     #   replace the footer key/value metadata instead of copying it
@@ -235,13 +244,14 @@ module Herringbone
     # @option writer_options [EncryptionConfiguration, Hash{Symbol => Object}, false] :encryption (as the input) see Writer;
     #   false for a plaintext output
     # @return [Report] what was done
-    # @raise [ArgumentError] when the redaction does not fit the file's schema, or for a writer
-    #   option that does not apply (+row_group_bytes:+, +row_group_rows:+)
+    # @raise [ArgumentError] when the redaction does not fit the file's schema, for a writer
+    #   option that does not apply (+row_group_bytes:+, +row_group_rows:+), or when a Reader comes
+    #   with +decryption:+
     # @raise [EncodeError] when a replacement value cannot be written to its column
     # @raise [DecryptionError] when the input is encrypted and a key it needs was not given
-    def apply(input_io, output_io, decryption: nil, **writer_options)
+    def apply(io_or_reader, output_io, decryption: nil, **writer_options)
       check_scopes!
-      Rewriter.new(self, input_io, decryption: decryption).apply(output_io, **writer_options)
+      Rewriter.new(self, io_or_reader, decryption: decryption).apply(output_io, **writer_options)
     end
 
     # Short summary for the console: one clause per statement, naming the columns but not the

data/lib/herringbone/restricted_readable_io.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module Herringbone
+  # Wraps every IO Herringbone reads a Parquet file from, and lets the reading code use only
+  # #read, #seek, #pos and #size. Any object implementing #read and #seek can then be read from
+  # (a File, a StringIO, an object fetching ranges from S3), and the reading code cannot start
+  # depending on #readpartial, #pread, #eof? or anything else only some IOs have: what it needs
+  # beyond the four methods it has to build on top of them.
+  #
+  # #path is there only to name the file in messages and in Inspector output.
+  #
+  #   reader = Herringbone::Reader.new(File.open("data.parquet", "rb"))
+  #   reader.io      # => #<Herringbone::RestrictedReadableIO data.parquet>
+  #   reader.io.path # => "data.parquet"
+  class RestrictedReadableIO
+    # What #path returns when the IO has no path of its own (a StringIO, a socket-backed IO...)
+    UNTITLED = "untitled.parquet"
+
+    # Wraps +io+, unless it is already wrapped
+    #
+    # @param io [IO, StringIO, RestrictedReadableIO, #read] anything responding to #read and #seek
+    # @return [RestrictedReadableIO] +io+ itself when it is one, a new wrapper of it otherwise
+    # @raise [ArgumentError] when +io+ does not respond to #read and #seek
+    def self.wrap(io)
+      io.is_a?(self) ? io : new(io)
+    end
+
+    # @param io [IO, StringIO, #read] anything responding to #read and #seek
+    # @raise [ArgumentError] when +io+ does not respond to #read and #seek
+    def initialize(io)
+      unless io.respond_to?(:read) && io.respond_to?(:seek)
+        raise ArgumentError, "#{self.class} wraps an IO that supports #seek and #read, got #{io.class}"
+      end
+      @io = io
+    end
+
+    # Reads up to +n_bytes+ from the current position
+    #
+    # @param n_bytes [Integer] bytes wanted
+    # @return [String, nil] at most +n_bytes+, fewer when the end of the IO comes first; nil at the end
+    def read(n_bytes)
+      @io.read(n_bytes)
+    end
+
+    # Moves to +offset+ bytes from the start of the IO
+    #
+    # @param offset [Integer] absolute position
+    # @return [Integer] 0, as IO#seek returns
+    def seek(offset)
+      @io.seek(offset)
+      0
+    end
+
+    # @return [Integer] the current position, in bytes from the start of the IO
+    def pos
+      @io.pos
+    end
+
+    # Size of the IO in bytes, asked of the IO when it can tell (File, StringIO, Tempfile...) and
+    # found by seeking to its end otherwise
+    #
+    # @return [Integer] the size in bytes
+    def size
+      return @io.size if @io.respond_to?(:size)
+      at = @io.pos
+      begin
+        @io.seek(0, IO::SEEK_END)
+        @io.pos
+      ensure
+        @io.seek(at)
+      end
+    end
+
+    # @return [String] the path of the IO (File#path, Tempfile#path), or UNTITLED when it has none
+    def path
+      path = @io.path if @io.respond_to?(:path)
+      path ? path.to_s : UNTITLED
+    end
+
+    # @return [String] e.g. "#<Herringbone::RestrictedReadableIO data.parquet>"
+    def inspect
+      "#<#{self.class.name} #{path}>"
+    end
+  end
+end