herringbone 0.6.1 → 0.8.0

data/lib/herringbone/codecs/snappy.rb

@@ -14,6 +14,8 @@ module Herringbone
     #
     # 32-bit loads are done with four getbyte calls rather than unpack1(offset:) to stay
     # compatible with Ruby 3.0 - the speed difference on MRI is marginal.
+    #
+    # @api private
     module Snappy
       # Raised for corrupt or truncated Snappy input, and for input too large to compress.
       class Error < StandardError; end
@@ -227,32 +229,72 @@ module Herringbone
         out
       end
 
-      # @param input [String] bytes to compress
+      # Several Strings are compressed as their concatenation without joining them: fragments are
+      # compressed straight out of the part holding them, and only a fragment straddling two parts
+      # is copied out. The output is the same as for the joined input.
+      #
+      # @param input [String, Array<String>] bytes to compress, or the parts of them
       # @return [String] raw snappy block in ASCII-8BIT
       # @raise [Error] if the input is larger than MAX_UNCOMPRESSED
       def compress(input)
-        src = (input.encoding == Encoding::BINARY) ? input : input.b
-        n = src.bytesize
+        parts = input.is_a?(Array) ? input.map { |part| binary(part) } : [binary(input)]
+        n = parts.sum(&:bytesize)
         # Checked for libsnappy too: it would truncate the 32-bit length preamble without complaint
         raise Error, "input too large for snappy" if n > MAX_UNCOMPRESSED
         if (lib = native)
-          return lib.deflate(src)
+          return lib.deflate((parts.size == 1) ? parts.first : parts.join)
         end
 
         out = String.new(capacity: 32 + n + n / 6, encoding: Encoding::BINARY)
         write_varint(out, n)
         table = Array.new(1 << HASH_BITS, 0)
+        part_index = 0
+        offset = 0
         start = 0
         while start < n
           len = n - start
           len = BLOCK_SIZE if len > BLOCK_SIZE
           table.fill(0)
-          compress_block(src, start, len, out, table)
+          part_index, offset = next_part(parts, part_index, offset)
+          part = parts[part_index]
+          if offset + len <= part.bytesize
+            compress_block(part, offset, len, out, table)
+            offset += len
+          else
+            fragment = String.new(capacity: len, encoding: Encoding::BINARY)
+            while fragment.bytesize < len
+              part_index, offset = next_part(parts, part_index, offset)
+              take = [len - fragment.bytesize, parts[part_index].bytesize - offset].min
+              fragment << parts[part_index].byteslice(offset, take)
+              offset += take
+            end
+            compress_block(fragment, 0, len, out, table)
+          end
           start += len
         end
         out
       end
 
+      # @param str [String] bytes in any encoding
+      # @return [String] +str+, or a binary copy when it is in another encoding
+      def binary(str)
+        (str.encoding == Encoding::BINARY) ? str : str.b
+      end
+
+      # Skips past parts that are used up
+      #
+      # @param parts [Array<String>] the input parts
+      # @param index [Integer] index of the current part
+      # @param offset [Integer] offset in the current part
+      # @return [Array(Integer, Integer)] index of a part with bytes left at the offset, and that offset
+      def next_part(parts, index, offset)
+        while offset >= parts[index].bytesize
+          index += 1
+          offset = 0
+        end
+        [index, offset]
+      end
+
       # Reads the uncompressed-length preamble (a little-endian base-128 varint) at the start of +src+.
       # @param src [String] raw snappy block
       # @param n [Integer] byte size of +src+
@@ -490,7 +532,7 @@ module Herringbone
         end
       end
 
-      private_class_method :native, :native_library, :decompress_io_buffer, :decompress_string, :read_varint, :write_varint, :compress_block, :block_words, :match_length_words,
+      private_class_method :native, :native_library, :binary, :next_part, :decompress_io_buffer, :decompress_string, :read_varint, :write_varint, :compress_block, :block_words, :match_length_words,
         :match_length, :emit_literal, :emit_copy, :emit_copy_upto64
     end
   end

data/lib/herringbone/combiner.rb

@@ -0,0 +1,426 @@
+# frozen_string_literal: true
+
+module Herringbone
+  # Concatenates Parquet files into one, behind Herringbone.combine. Building it reads the footer of
+  # every input and checks every input against the output schema, so all mismatches surface at
+  # once, before anything is written.
+  #
+  # Every row group of every input becomes a row group of the output, in order. Column chunks are
+  # copied byte for byte wherever the output stores a column's values as the input does, and only
+  # their offsets (in the column metadata, the page index and the bloom filter references) are
+  # rebased. That holds per leaf column, so a narrower input still has most of its chunks copied.
+  # Columns that cannot be copied are encoded again from their values: those the output widens in
+  # physical type, time unit or nullability, and those encrypted in the input or to be encrypted in
+  # the output. Fields an input lacks are written as nulls.
+  #
+  # @api private
+  class Combiner
+    # What #apply did
+    #
+    # @!attribute rows
+    #   @return [Integer] rows written
+    # @!attribute row_groups
+    #   @return [Hash{Symbol => Integer}] +{copied:, rewritten:}+ row group counts: a row group is
+    #     rewritten when at least one of its column chunks had to be encoded again
+    # @!attribute inputs
+    #   @return [Array<InputReport>] how each input was fitted to the output schema, in input order
+    # @api public
+    Report = Struct.new(:rows, :row_groups, :inputs, keyword_init: true)
+
+    # How one input was fitted to the output schema. Paths are dotted field paths ("address.zip").
+    #
+    # @!attribute name
+    #   @return [String] "input 2", with the IO's path when it has one: "input 2 (2026-03.parquet)"
+    # @!attribute filled
+    #   @return [Array<String>] fields the input lacks, written as nulls
+    # @!attribute widened
+    #   @return [Array<String>] fields the output holds in a wider type, or nullable
+    # @!attribute dropped
+    #   @return [Array<String>] fields of the input the output does not have (+schema: :intersect+)
+    # @api public
+    InputReport = Struct.new(:name, :filled, :widened, :dropped, keyword_init: true)
+
+    # One input file: its Reader (with the combiner's own options), its ChunkCopier, output column
+    # index => the input's column holding its values (output columns the input lacks are not in
+    # it), and its InputReport
+    Input = Struct.new(:reader, :copier, :sources, :report)
+
+    # Closing paragraph of the message when inputs do not fit the output schema
+    FIT_ADVICE = <<~ADVICE.chomp
+      Pass a schema every input fits (schema: :union makes one), or schema: :intersect to keep only
+      the fields all inputs have.
+    ADVICE
+
+    # Encryption settings the inherited encryption takes from the inputs => how messages name them
+    ENCRYPTION_SETTINGS = {
+      footer_key: "footer key", footer_key_metadata: "footer key metadata", plaintext_footer: "footer mode",
+      algorithm: "algorithm", aad_prefix: "AAD prefix", store_aad_prefix: "AAD prefix storage"
+    }.freeze
+    # Settings whose values messages leave out
+    SECRET_SETTINGS = %i[footer_key aad_prefix].freeze
+
+    # Converted annotations of time and timestamp columns => their unit
+    CONVERTED_UNITS = {
+      Format::ConvertedType::TIME_MILLIS => :millis, Format::ConvertedType::TIME_MICROS => :micros,
+      Format::ConvertedType::TIMESTAMP_MILLIS => :millis, Format::ConvertedType::TIMESTAMP_MICROS => :micros
+    }.freeze
+
+    # @return [Schema] the schema of the output
+    attr_reader :schema
+
+    # @param ios_or_readers [Enumerable<IO, StringIO, Reader>] the input files: IOs read with #seek
+    #   and #read, or Readers (which bring their own +decryption:+); enumerated once, and none is
+    #   closed
+    # @param schema [Schema, Symbol, nil] the output schema: nil when every input has the same one,
+    #   +:union+ for the union of the inputs' schemas, +:intersect+ for their intersection, or a
+    #   Schema each input must fit
+    # @param decryption [DecryptionConfiguration, Hash{Symbol => Object}, Array, #call, nil] keys of
+    #   encrypted inputs given as IOs, see Reader.new
+    # @raise [ArgumentError] when there are no inputs, an input is neither an IO nor a Reader, or
+    #   +schema:+ is not one of the above
+    # @raise [IncompatibleSchema] listing every input (and every field of it) that does not fit
+    # @raise [FormatError] when an input's footer cannot be read
+    # @raise [DecryptionError] when an input's footer is encrypted and cannot be decrypted
+    def initialize(ios_or_readers, schema: nil, decryption: nil)
+      unless schema.is_a?(Schema) || [nil, :union, :intersect].include?(schema)
+        raise ArgumentError, "schema: takes :union, :intersect or a Herringbone::Schema, got #{schema.inspect}"
+      end
+      readers = open_readers(ios_or_readers, decryption)
+      @names = readers.each_with_index.map do |reader, k|
+        path = reader.io.path
+        path.equal?(RestrictedReadableIO::UNTITLED) ? "input #{k}" : "input #{k} (#{path})"
+      end
+      @conflicts = []
+      @schema = output_schema(schema, readers)
+      @inputs = readers.each_with_index.map { |reader, k| fit(reader, k, schema == :intersect) }
+      return if @conflicts.empty?
+      what = {union: "the union of the inputs", intersect: "the intersection of the inputs"}.fetch(schema, "given as schema:")
+      raise_conflicts("#{inputs_count("does", "do")} not fit the schema (#{what}):", "the schema", advice: FIT_ADVICE)
+    end
+
+    # @param output_io [IO, #write] destination, written sequentially; not closed
+    # @param options [Hash{Symbol => Object}] Writer options for chunks encoded again, and
+    #   +metadata:+ / +encryption:+ for the output
+    # @option options [Symbol] :compression (codec of each source chunk) codec for chunks encoded
+    #   again; written nulls use the Writer's default
+    # @option options [Hash{String => String}] :metadata (the first input's) footer key/value metadata
+    # @option options [EncryptionConfiguration, Hash{Symbol => Object}, Key, String, false] :encryption
+    #   (encrypted like the encrypted inputs, if any) see Writer; false for a plaintext output
+    # @option options [Integer, nil] :compression_level (nil) level for that codec, see Writer
+    # @option options [Boolean, Array<String>, Hash{String => Boolean, Hash}] :bloom_filters (nil)
+    #   columns whose encoded chunks get a bloom filter, besides those whose source chunk had one
+    # @option options [Integer] :page_bytes (1MB) approximate uncompressed data page size
+    # @option options [Integer] :page_rows (20_000) maximum rows per data page
+    # @option options [Integer] :data_page_version (1) 1 or 2
+    # @option options [Boolean, Array<String>] :dictionary (true) see Writer
+    # @option options [Hash{String => Symbol}] :encodings ({}) see Writer
+    # @return [Report]
+    # @raise [ArgumentError] for +row_group_bytes:+ / +row_group_rows:+, an invalid writer option,
+    #   or encrypted inputs that are encrypted differently while +encryption:+ is not given
+    # @raise [DecryptionError] when a column is encrypted and its key was not given
+    def apply(output_io, **options)
+      options = {metadata: @inputs.first.copier.metadata(@schema)}.merge(options)
+      options[:encryption] = inherited_encryption unless options.key?(:encryption)
+      @keep_codecs = !options.key?(:compression)
+      report = Report.new(rows: 0, row_groups: {copied: 0, rewritten: 0}, inputs: @inputs.map(&:report))
+      Writer.open_for_copies(output_io, @schema, "combining keeps the row groups of the inputs", **options) do |writer|
+        @inputs.each do |input|
+          from = 0
+          input.reader.row_groups.each_with_index do |rg, i|
+            report.row_groups[write_row_group(input, i, from, writer)] += 1 unless rg.num_rows.zero?
+            from += rg.num_rows
+          end
+        end
+        report.rows = writer.rows_written
+      end
+      report
+    end
+
+    private
+
+    # Readers with the combiner's own options. A Reader given as input lends its IO and its
+    # decryption, but not its +keys:+ or +time_zone:+, which would change the values encoded again.
+    #
+    # @param ios_or_readers [Enumerable<IO, StringIO, Reader>] the inputs
+    # @param decryption [Object, nil] keys for the inputs given as IOs
+    # @return [Array<Reader>]
+    # @raise [ArgumentError] when there are no inputs, or one is neither an IO nor a Reader
+    def open_readers(ios_or_readers, decryption)
+      # A single IO is Enumerable too (over its lines), hence the check for #read
+      if ios_or_readers.respond_to?(:read) || !ios_or_readers.respond_to?(:each)
+        raise ArgumentError, "Herringbone.combine takes an Enumerable of IOs or Herringbone::Readers, " \
+          "got #{ios_or_readers.class}: Herringbone.combine([a, b], output_io)"
+      end
+      readers = ios_or_readers.each_with_index.map do |item, k|
+        next Reader.new(item.io, decryption: item.decryption) if item.is_a?(Reader)
+        next Reader.new(item, decryption: decryption) if item.respond_to?(:seek) && item.respond_to?(:read)
+        hint = "; Herringbone does not open files by path, pass File.open(path, \"rb\")" if item.is_a?(String) || item.respond_to?(:to_path)
+        raise ArgumentError, "Input #{k} is a #{item.class}, not an IO or a Herringbone::Reader#{hint}"
+      end
+      raise ArgumentError, "Herringbone.combine needs at least one input" if readers.empty?
+      readers
+    end
+
+    # @param schema [Schema, Symbol, nil] the +schema:+ option
+    # @param readers [Array<Reader>] the inputs
+    # @return [Schema]
+    # @raise [IncompatibleSchema] when the inputs' schemas differ (nil), or do not unite or intersect
+    def output_schema(schema, readers)
+      return schema if schema.is_a?(Schema)
+      first = readers.first.schema
+      readers.each_with_index.drop(1).each do |reader, k|
+        if schema.nil?
+          next if reader.schema == first
+          before = @conflicts.size
+          differences(first, reader.schema, k, "input 0")
+          if @conflicts.size == before
+            reordered = reader.schema.fields.map(&:name) != first.fields.map(&:name)
+            conflict([], nil, nil, reordered ? "the fields are in another order than in input 0" :
+              "the same fields as input 0, annotated or laid out another way", k)
+          end
+        else
+          first = (schema == :union) ? first + reader.schema : first & reader.schema
+        end
+      rescue IncompatibleSchema => e
+        e.conflicts.each { |c| conflict(c.path, c.left, c.right, c.reason, k) }
+      end
+      return first if @conflicts.empty?
+      if schema.nil?
+        raise_conflicts("#{inputs_count("has", "have")} another schema than #{@names[0]}:", "input 0", advice: <<~ADVICE.chomp)
+          Combine files with the same schema, or pass schema: :union to fill the fields an input lacks
+          with nulls and widen the others, or schema: :intersect to keep only the fields all inputs have.
+        ADVICE
+      end
+      raise_conflicts("Cannot #{(schema == :union) ? "unite" : "intersect"} the schemas of the inputs, " \
+        "#{inputs_count("does", "do")} not fit those before it:", "the inputs before it")
+    end
+
+    # Each field of +mine+ with the field of the same name in +theirs+ (nil when there is none),
+    # then the same for their members down through structs, lists and maps
+    #
+    # @param mine [Schema] the schema walked
+    # @param theirs [Schema] the schema its fields are looked up in
+    # @return [Array<Array(Schema::Field, Schema::Field, String)>] field, twin and dotted path
+    def pairs(mine, theirs)
+      walk = lambda do |a, b, path|
+        [[a, b, path.join(".")]] + case (b && a.kind == b.kind) ? a.kind : nil
+        when :struct
+          by_name = b.children_by_name
+          a.children.flat_map { |c| walk.call(c, by_name[c.name], path + [c.name]) }
+        when :list then walk.call(a.element, b.element, path + ["element"])
+        when :map then walk.call(a.key, b.key, path + ["key"]) + ((a.value && b.value) ? walk.call(a.value, b.value, path + ["value"]) : [])
+        else []
+        end
+      end
+      mine.fields.flat_map { |f| walk.call(f, theirs.field(f.name), [f.name]) }
+    end
+
+    # Where the fields of +theirs+ differ from those of +mine+, recorded as conflicts
+    #
+    # @param mine [Schema] the schema compared with
+    # @param theirs [Schema] the input's schema
+    # @param k [Integer] the input's position
+    # @param other [String] what to call +mine+, e.g. "input 0"
+    # @param skip [Array<String>] paths not to compare
+    # @return [void]
+    def differences(mine, theirs, k, other, skip: [])
+      pairs(mine, theirs).each do |a, b, path|
+        next if skip.include?(path)
+        if b.nil? then conflict(path, nil, nil, "missing here, #{other} has it", k)
+        elsif (da = describe(a)) != (db = describe(b)) then conflict(path, da, db, nil, k)
+        elsif a.optional != b.optional
+          conflict(path, a.optional ? "nullable" : "required", b.optional ? "nullable" : "required", nil, k)
+        elsif a.node.field_id != b.node.field_id
+          conflict(path, "field_id #{a.node.field_id.inspect}", "field_id #{b.node.field_id.inspect}", nil, k)
+        end
+      end
+      pairs(theirs, mine).each { |_, a, path| conflict(path, nil, nil, "only here, #{other} lacks it", k) unless a }
+    end
+
+    # Fits an input to the output schema, recording the conflicts when it does not fit. The input
+    # must be the same as the output or narrower: uniting it with the output (Schema#+) must leave
+    # the output as it is. Each output leaf is found in the input through the logical tree, by
+    # name, so the element of a pyarrow list ("list.item") or of a legacy 2-level list is still the
+    # same column.
+    #
+    # @param reader [Reader] the input
+    # @param k [Integer] its position
+    # @param drop [Boolean] whether fields the output lacks are dropped (+schema: :intersect+)
+    #   rather than refused
+    # @return [Input]
+    def fit(reader, k, drop)
+      theirs = reader.schema
+      pairs = pairs(@schema, theirs)
+      filled = pairs.filter_map { |_, b, path| path unless b }
+      dropped = pairs(theirs, @schema).filter_map { |_, a, path| path unless a }
+      if pairs.all? { |_, b, _| b.nil? }
+        conflict([], @schema.fields.map(&:name).join(", "), theirs.fields.map(&:name).join(", "), "no fields in common with the schema", k)
+      else
+        pairs.each { |a, b, path| conflict(path, nil, nil, "missing here, and the schema requires it", k) if b.nil? && !a.optional }
+        begin
+          fitted = @schema + (drop ? theirs & @schema : theirs)
+          differences(@schema, fitted, k, "the schema", skip: filled)
+        rescue IncompatibleSchema => e
+          e.conflicts.each { |c| conflict(c.path, c.left, c.right, c.reason, k) }
+        end
+      end
+      sources = pairs.filter_map { |a, b, _| [a.column.index, b.column] if b && a.leaf? && b.leaf? }.to_h
+      widened = pairs.filter_map { |a, b, path| path if b && (a.optional != b.optional || (a.leaf? && describe(a) != describe(b))) }
+      report = InputReport.new(name: @names[k], filled: filled, widened: widened, dropped: dropped)
+      Input.new(reader, Reader::ChunkCopier.new(reader), sources, report)
+    end
+
+    # Whether a chunk of +theirs+ holds the values of +mine+ byte for byte. For a column of an
+    # input that fits the output, the widenings that keep the bytes are those of the annotation
+    # alone: int8 into int32, uint16 into int32, string into binary... Equal levels mean that no
+    # field on the way became nullable, since nullability only ever widens here.
+    #
+    # @param mine [Schema::Column] a column of the output schema
+    # @param theirs [Schema::Column] the input's column for it
+    # @return [Boolean]
+    def same_bytes?(mine, theirs)
+      mine.type == theirs.type && mine.type_length == theirs.type_length &&
+        mine.max_definition_level == theirs.max_definition_level &&
+        mine.max_repetition_level == theirs.max_repetition_level &&
+        time_unit(mine.node) == time_unit(theirs.node)
+    end
+
+    # @param node [Schema::Node] a leaf
+    # @return [Symbol, nil] the unit of a time or timestamp column, nil for other columns
+    def time_unit(node)
+      lt = node.logical_type
+      (lt&.timestamp || lt&.time)&.unit&.to_sym || CONVERTED_UNITS[node.converted_type]
+    end
+
+    # @param field [Schema::Field] any field
+    # @return [String] its type as messages name it, see Schema::Merge#describe
+    def describe(field) = (@describer ||= Schema::Merge.new(:union)).describe(field)
+
+    # @param path [String, Array<String>] path of the field, empty for the input as a whole
+    # @param left [String, nil] the field as the schema compared with has it
+    # @param right [String, nil] the field as the input has it
+    # @param reason [String, nil] why they do not fit
+    # @param k [Integer] the input's position
+    # @return [void]
+    def conflict(path, left, right, reason, k)
+      @conflicts << IncompatibleSchema::Conflict.new(Array(path).join("."), left, right, reason, k)
+    end
+
+    # @param singular [String] the verb for one input, e.g. "does"
+    # @param plural [String] the verb for several, e.g. "do"
+    # @return [String] the inputs with conflicts and the verb: "2 of 5 inputs do", "1 of 5 inputs
+    #   does", or "The input does" when there is only one
+    def inputs_count(singular, plural)
+      n = @conflicts.map(&:input).uniq.size
+      return "The input #{singular}" if @names.size == 1
+      "#{n} of #{@names.size} inputs #{(n == 1) ? singular : plural}"
+    end
+
+    # Raises the conflicts, grouped by input, as one IncompatibleSchema:
+    #
+    #   2 of 3 inputs do not fit the schema (schema:):
+    #
+    #     input 1 (2026-02.parquet)
+    #       price        double in the schema, string here (no common type)
+    #       address.zip  required in the schema, nullable here
+    #
+    # @param header [String] the first line
+    # @param other [String] what to call the schema the inputs are compared with, e.g. "the schema"
+    # @param advice [String, nil] a closing paragraph
+    # @return [void]
+    # @raise [IncompatibleSchema] always
+    def raise_conflicts(header, other, advice: nil)
+      groups = @conflicts.group_by(&:input).map do |k, conflicts|
+        width = conflicts.map { |c| c.path.size }.max
+        lines = conflicts.map do |c|
+          if c.path.empty? && c.left
+            "#{c.reason}\n      #{other}: #{c.left.delete_prefix("fields ")}\n      here: #{c.right.delete_prefix("fields ")}"
+          elsif c.path.empty?
+            c.reason
+          else
+            detail = [("#{c.left} in #{other}, #{c.right} here" if c.left), ("(#{c.reason})" if c.reason && c.left), (c.reason unless c.left)]
+            "#{c.path.ljust(width)}  #{detail.compact.join(" ")}"
+          end
+        end
+        "  #{@names[k]}\n#{lines.map { |l| "    #{l}" }.join("\n")}"
+      end
+      message = [header, groups.join("\n\n"), advice].compact.join("\n\n")
+      raise IncompatibleSchema.new(@conflicts, message: message)
+    end
+
+    # The +encryption:+ option that encrypts the output like the encrypted inputs: the same
+    # algorithm, footer mode, footer key and key metadata, AAD prefix, and the same key for each
+    # column. Plaintext inputs get that encryption too.
+    #
+    # @return [EncryptionConfiguration, nil] nil when no input is encrypted
+    # @raise [ArgumentError] listing how the encrypted inputs differ, when they do
+    # @raise [DecryptionError] when a key the output needs was not given
+    def inherited_encryption
+      encrypted = @inputs.select { |input| input.reader.decryptor }
+      return nil if encrypted.empty?
+      settings = encrypted.map { |input| input.reader.decryptor.writer_settings(nil).to_h.except(:columns) }
+      problems = settings.first.keys.filter_map do |setting|
+        groups = encrypted.zip(settings).group_by { |_, h| h[setting] }
+        next if groups.size == 1
+        "#{ENCRYPTION_SETTINGS.fetch(setting)} differs: " + groups.map { |value, members|
+          "#{members.map { |input, _| input.report.name }.join(", ")}#{" (#{value.inspect})" unless SECRET_SETTINGS.include?(setting)}"
+        }.join(" / ")
+      end
+      columns = {}
+      owners = {}
+      layouts = encrypted.map { |input| input.copier.encrypted_columns(input.sources.to_h { |j, col| [@schema.columns[j].dotted_path, col] }) }
+      encrypted.zip(layouts).each do |input, layout|
+        layout.each do |path, setting|
+          owner = owners[path] ||= input.report.name
+          columns[path] ||= setting
+          problems << "#{path} is encrypted with another key in #{input.report.name} than in #{owner}" if columns[path] != setting
+        end
+      end
+      unless problems.empty?
+        raise ArgumentError, <<~MESSAGE.chomp
+          The encrypted inputs are encrypted differently, so there is no "encrypted like the inputs":
+
+          #{problems.map { |p| "  #{p}" }.join("\n")}
+
+          Pass encryption: for the output (a Herringbone::Key or an EncryptionConfiguration), or
+          encryption: false to write it in plaintext.
+        MESSAGE
+      end
+      uniform = encrypted.zip(layouts).all? { |input, layout| layout.size == input.sources.size && layout.values.all?(:footer) }
+      EncryptionConfiguration.new(**settings.first, columns: uniform ? nil : columns)
+    end
+
+    # Writes row group +i+ of +input+, copying the column chunks it can. A field with a column that
+    # cannot be copied is read as a whole, since Reader and Writer handle top-level fields; the
+    # columns of it that can be copied still are.
+    #
+    # @param input [Input] the input file
+    # @param i [Integer] row group index in the input
+    # @param from [Integer] index of the row group's first row in the input
+    # @param writer [Writer] the output
+    # @return [Symbol] +:copied+ or +:rewritten+
+    def write_row_group(input, i, from, writer)
+      chunks = input.reader.row_groups[i].columns
+      n = input.reader.row_groups[i].num_rows
+      copies = {}
+      codecs = {}
+      blooms = {}
+      input.sources.each do |j, source|
+        if same_bytes?(@schema.columns[j], source) && !writer.encrypted_column?(j) && !chunks.fetch(source.index).crypto_metadata
+          copies[j] = input.copier.copy(i, source)
+        else
+          codec, bloom = input.copier.recode_settings(i, source)
+          codecs[j] = codec if codec && @keep_codecs
+          blooms[j] = true if bloom
+        end
+      end
+      encode = @schema.fields.reject { |field| field.leaves.all? { |col| copies.key?(col.index) } }
+      present = encode.map(&:name) & input.reader.schema.fields.map(&:name)
+      data = present.empty? ? {} : input.reader.read(as: :columns, columns: present, from: from, limit: n)
+      encode.each { |field| writer.buffer_field(field.name, data.fetch(field.name) { Array.new(n) }) }
+      sorting = input.copier.sorting_columns(i, input.sources.to_h { |j, source| [source.index, j] })
+      writer.write_row_group(n, copies: copies, codecs: codecs, bloom_filters: blooms, sorting_columns: sorting)
+      (copies.size == @schema.columns.size) ? :copied : :rewritten
+    end
+  end
+end

data/lib/herringbone/compression.rb

@@ -4,7 +4,8 @@ require "zlib"
 require "stringio"
 
 module Herringbone
-  # Raised when a file uses (or a writer asks for) a codec whose library is not loaded
+  # Raised when a file uses (or a writer asks for) a codec whose library is not loaded. The message
+  # says which gem to install, and that it has to be required before the read or write.
   class MissingCodecError < UnsupportedError
     # @return [String] codec name (+"ZSTD"+) and name of the gem providing it (+"zstd-ruby"+)
     attr_reader :codec, :gem_name
@@ -12,11 +13,29 @@ module Herringbone
     # @param codec [String] codec name, as shown in the message
     # @param gem_name [String] gem to add to the Gemfile
     # @param path [String] what to require to load the gem
-    def initialize(codec, gem_name, path)
+    # @param column [String, nil] dotted path of the column being read, when there is one
+    def initialize(codec, gem_name, path, column: nil)
       @codec = codec
       @gem_name = gem_name
-      super("#{codec} compression needs the \"#{gem_name}\" gem, which is not loaded. Add " \
-        "`gem \"#{gem_name}\"` to your Gemfile and `require \"#{path}\"` to use #{codec}.")
+      @path = path
+      where = column ? " to decompress column #{column}" : ""
+      super(<<~MESSAGE.chomp)
+        #{codec} compression needs the "#{gem_name}" gem#{where}, but it is not loaded.
+        Herringbone does not load optional codec gems by itself. To use #{codec}:
+
+          1. Install the gem: add `gem "#{gem_name}"` to your Gemfile, or run `gem install #{gem_name}`.
+          2. Load it with `require "#{path}"` before the read or write that needs #{codec},
+             for instance at the top of your script or in an initializer.
+      MESSAGE
+    end
+
+    # The same error, naming the column that could not be decompressed
+    #
+    # @param column [String] dotted path of the column
+    # @return [MissingCodecError]
+    # @api private
+    def for_column(column)
+      self.class.new(@codec, @gem_name, @path, column: column).tap { |e| e.set_backtrace(backtrace) if backtrace }
     end
   end
 
@@ -25,6 +44,8 @@ module Herringbone
   # come from optional gems (zstd-ruby, brotli), which the application requires; if they are
   # not loaded, MissingCodecError says what to add.
   # (Snappy also uses the optional snappy gem when it is loaded, see Codecs::Snappy.)
+  #
+  # @api private
   module Compression
     module_function
 
@@ -129,7 +150,7 @@ module Herringbone
     def decompress(codec, data, uncompressed_size)
       return "".b if uncompressed_size.zero? && data.empty?
       out = case codec
-      when Format::Codec::UNCOMPRESSED then data
+      when Format::Codec::UNCOMPRESSED then (data.encoding == Encoding::BINARY) ? data : data.b
       when Format::Codec::SNAPPY then Codecs::Snappy.decompress(data)
       when Format::Codec::GZIP then gunzip(data)
       when Format::Codec::LZ4_RAW then Codecs::LZ4.decompress_block(data, uncompressed_size)
@@ -140,33 +161,75 @@ module Herringbone
       else
         raise UnsupportedError, "Unsupported compression codec #{Format::Codec::NAMES[codec] || codec}"
       end
-      out = out.b
+      out = binary(out)
       if out.bytesize != uncompressed_size
         raise FormatError, "Decompressed #{out.bytesize} bytes, expected #{uncompressed_size}"
       end
       out
     end
 
-    # Compresses a page body
+    # Compresses a page body. Given several Strings, compresses their concatenation; GZIP, ZSTD
+    # and Snappy take them one by one, the other codecs join them first.
     #
     # @param codec [Integer] Format::Codec id
-    # @param data [String] bytes to compress
+    # @param data [String, Array<String>] bytes to compress, or the parts of them
     # @param level [Integer, nil] compression level checked with check_level!, nil for the default
     # @return [String] compressed bytes (binary)
     # @raise [UnsupportedError] for a codec herringbone does not implement
     # @raise [MissingCodecError] when the codec's gem is not loaded
     def compress(codec, data, level = nil)
+      parts = data.is_a?(Array) ? data : [data]
       case codec
-      when Format::Codec::UNCOMPRESSED then data
-      when Format::Codec::SNAPPY then Codecs::Snappy.compress(data)
-      when Format::Codec::GZIP then level ? Zlib.gzip(data, level: level) : Zlib.gzip(data)
+      when Format::Codec::GZIP then return binary(gzip(parts, level))
+      when Format::Codec::ZSTD then return binary(zstd(parts, level))
+      when Format::Codec::SNAPPY then return binary(Codecs::Snappy.compress(data))
+      end
+      data = parts.join if data.is_a?(Array)
+      out = case codec
+      when Format::Codec::UNCOMPRESSED then (data.encoding == Encoding::BINARY) ? data : data.b
       when Format::Codec::LZ4_RAW then Codecs::LZ4.compress_block(data)
       when Format::Codec::LZ4 then Codecs::LZ4.compress_hadoop(data)
-      when Format::Codec::ZSTD then level ? library(codec).compress(data, level: level) : library(codec).compress(data)
       when Format::Codec::BROTLI then level ? library(codec).deflate(data, quality: level) : library(codec).deflate(data)
       else
         raise UnsupportedError, "Unsupported compression codec #{Format::Codec::NAMES[codec] || codec}"
-      end.b
+      end
+      binary(out)
+    end
+
+    # A gzip member with no file name and a zero mtime, like Zlib.gzip makes
+    #
+    # @param parts [Array<String>] bytes to compress, in order
+    # @param level [Integer, nil] zlib level, nil for the default
+    # @return [String] gzip data
+    def gzip(parts, level)
+      z = Zlib::Deflate.new(level || Zlib::DEFAULT_COMPRESSION, Zlib::MAX_WBITS + 16)
+      parts.each { |part| z << part }
+      z.finish
+    ensure
+      z&.close
+    end
+
+    # @param parts [Array<String>] bytes to compress, in order
+    # @param level [Integer, nil] zstd level, nil for the default
+    # @return [String] a single zstd frame
+    # @raise [MissingCodecError] when zstd-ruby is not loaded
+    def zstd(parts, level)
+      zstd = library(Format::Codec::ZSTD)
+      if parts.size == 1
+        return level ? zstd.compress(parts.first, level: level) : zstd.compress(parts.first)
+      end
+      stream = level ? zstd::StreamingCompress.new(level: level) : zstd::StreamingCompress.new
+      parts.each { |part| stream << part }
+      stream.finish
+    end
+
+    # Codec output relabelled as binary without copying it, unless it is frozen
+    #
+    # @param str [String] bytes
+    # @return [String] +str+ or a binary copy of it
+    def binary(str)
+      return str if str.encoding == Encoding::BINARY
+      str.frozen? ? str.b : str.force_encoding(Encoding::BINARY)
     end
 
     # Handles files whose gzip data consists of several concatenated members