herringbone 0.1.0 → 0.2.0

data/lib/herringbone/writer.rb

@@ -8,27 +8,37 @@ module Herringbone
   #     string :name
   #     list :tags, :string
   #   end
-  #   Herringbone::Writer.open("out.parquet", schema) do |w|
-  #     w << { "id" => 1, "name" => "one", "tags" => ["a", "b"] }
-  #     w << [2, "two", []]           # Arrays are taken in schema order
-  #     w << order                    # objects responding to #attributes (ActiveRecord) or #to_h
+  #   File.open("out.parquet", "wb") do |file|
+  #     Herringbone::Writer.open(file, schema) do |w|
+  #       w << { "id" => 1, "name" => "one", "tags" => ["a", "b"] }
+  #       w << [2, "two", []]           # Arrays are taken in schema order
+  #       w << order                    # objects responding to #attributes (ActiveRecord) or #to_h
+  #     end
   #   end
   #
-  # +schema+ is a Herringbone::Schema or a Hash spec (see Schema.define). The target is a path or an IO;
-  # paths are written to a temporary file next to the target and renamed into place on close, so a
-  # failed write never leaves a truncated file behind.
+  # The output is any IO that responds to #write (a File, StringIO, Tempfile, socket, pipe...); it
+  # is written sequentially and never seeked, rewound or closed by the writer. Herringbone does
+  # not open files by path.
   #
   # Options:
-  #   compression:     :zstd (default), :snappy, :gzip, :lz4 (LZ4_RAW), :brotli, :none
+  #   compression:     :snappy (default), :zstd, :gzip, :lz4 (LZ4_RAW), :lz4_hadoop, :brotli, :none
+  #                    (:zstd and :brotli need the zstd-ruby / brotli gems)
   #   row_group_bytes: flush a row group once the buffered values take roughly this much memory
   #                    (default 16MB). This bounds memory use while writing.
-  #   row_group_size:  also flush after this many rows (default: no row limit)
-  #   page_size:       approximate uncompressed data page size in bytes (default 1MB)
+  #   row_group_rows:  also flush after this many rows (default: no row limit)
+  #   page_bytes:      approximate uncompressed data page size (default 1MB)
+  #   page_rows:       at most this many rows per data page (default 20_000), which keeps the page
+  #                    index selective
   #   data_page_version: 1 (default) or 2
   #   dictionary:      true/false, or an Array of column paths to dictionary-encode
   #   encodings:       { "path.to.column" => :delta_binary_packed, ... } for non-dictionary pages
-  #   statistics:      write min/max/null_count statistics (default true)
   #   metadata:        Hash of String => String key/value metadata for the footer
+  #   bloom_filters:   write split block bloom filters: true (every column that supports them),
+  #                    an Array of column paths, or { "path" => true | { ndv:, fpp:, max_bytes: } }.
+  #                    Without ndv: the distinct values of each row group are counted. fpp defaults
+  #                    to 0.01 and max_bytes to 1MB. Filters are written after each row group.
+  #
+  # Statistics and page indexes (ColumnIndex/OffsetIndex) are always written.
   class Writer
     MAGIC = "PAR1"
     T = Format::Type
@@ -55,14 +65,15 @@ module Herringbone
 
     attr_reader :schema
 
-    # Opens a writer on a path or an IO. With a block, the file is finished when the block returns
-    # (or discarded if it raises) and the block's value is returned; without one, call #close.
-    def self.open(target, schema, **options)
-      writer = new(target, schema, **options)
+    # Opens a writer on +io+. With a block, the file is finished (footer written) when the block
+    # returns and the block's value is returned; if the block raises, the writer is aborted and no
+    # footer is written. Without a block, call #close to finish. The IO is never closed.
+    def self.open(io, schema, **options)
+      writer = new(io, schema, **options)
       return writer unless block_given?
       begin
         result = yield writer
-      rescue Exception # rubocop:disable Lint/RescueException -- also discard on Interrupt
+      rescue Exception # rubocop:disable Lint/RescueException -- also abort on Interrupt
         writer.abort
         raise
       end
@@ -70,29 +81,36 @@ module Herringbone
       result
     end
 
-    def initialize(target, schema, compression: :zstd, row_group_bytes: 16 * 1024 * 1024, row_group_size: nil,
-      page_size: 1024 * 1024, data_page_version: 1, dictionary: true, encodings: {}, statistics: true, metadata: {})
-      @schema = Schema.coerce(schema)
-      schema = @schema
+    def initialize(io, schema, compression: :snappy, row_group_bytes: 16 * 1024 * 1024, row_group_rows: nil,
+      page_bytes: 1024 * 1024, page_rows: 20_000, data_page_version: 1, dictionary: true, encodings: {},
+      metadata: {}, bloom_filters: nil)
+      raise ArgumentError, "Expected a Herringbone::Schema, got #{schema.class}" unless schema.is_a?(Schema)
+      @schema = schema
       @codec = Compression.codec_id(compression)
+      # Fail before creating any file if the codec's library is missing
+      Compression.ensure_available!(@codec)
       @row_group_bytes = Integer(row_group_bytes)
-      @row_group_size = row_group_size && Integer(row_group_size)
-      @row_limit = @row_group_size || ESTIMATE_AFTER_ROWS
-      @page_size = Integer(page_size)
+      @row_group_rows = row_group_rows && Integer(row_group_rows)
+      @row_limit = @row_group_rows || ESTIMATE_AFTER_ROWS
+      @page_bytes = Integer(page_bytes)
+      @page_rows = Integer(page_rows)
+      raise ArgumentError, "page_rows must be positive" unless @page_rows.positive?
+      @page_indexes = [] # [ColumnChunk, ColumnIndex or nil, OffsetIndex] per written column chunk
       @data_page_version = Integer(data_page_version)
       raise ArgumentError, "data_page_version must be 1 or 2" unless [1, 2].include?(@data_page_version)
       @dictionary = dictionary
       @encodings = encodings.to_h { |path, enc| [path.to_s, encoding_id(path, enc)] }
       unknown = @encodings.keys - schema.columns.map(&:dotted_path)
       raise ArgumentError, "encodings: no such column #{unknown.join(", ")}" unless unknown.empty?
-      @statistics = statistics
       @metadata = metadata
+      @bloom_filters = bloom_filter_config(bloom_filters)
+      @pending_bloom_filters = [] # [ColumnMetaData, BloomFilter] for the row group being written
       @row_groups = []
       @total_rows = 0
       @pos = 0
       @closed = false
       @bytes_per_row = nil
-      open_target(target)
+      @io = check_io!(io)
       write_raw(MAGIC)
       reset_buffers
     end
@@ -131,12 +149,6 @@ module Herringbone
       check_row_group_size if @buffered_rows >= @row_limit
       self
     end
-    alias_method :write, :<<
-
-    def write_rows(rows)
-      rows.each { |row| self << row }
-      self
-    end
 
     # Number of rows written so far, including buffered ones
     def rows_written = @total_rows + @buffered_rows
@@ -154,6 +166,7 @@ module Herringbone
         total_compressed_size: @pos - start,
         ordinal: @row_groups.size
       )
+      write_bloom_filters
       @total_rows += @buffered_rows
       reset_buffers
     end
@@ -161,13 +174,14 @@ module Herringbone
     def close
       return if @closed
       flush_row_group
+      write_page_indexes
       meta = Format::FileMetaData.new(
         version: 2,
         schema: @schema.to_elements,
         num_rows: @total_rows,
         row_groups: @row_groups,
         key_value_metadata: @metadata.empty? ? nil : @metadata.map { |k, v| Format::KeyValue.new(key: k.to_s, value: v&.to_s) },
-        created_by: "herringbone version #{VERSION}",
+        created_by: "herringbone-ruby #{VERSION}",
         column_orders: @schema.columns.map { Format::ColumnOrder.new(type_order: Format::TypeDefinedOrder.new) }
       )
       footer = meta.encode
@@ -176,20 +190,23 @@ module Herringbone
       write_raw(MAGIC)
       @io.flush if @io.respond_to?(:flush)
       @closed = true
-      finish_target
     end
 
-    # Stops writing without producing a file: a temporary file is deleted, an IO is left as is
+    # Stops writing without finishing the file (no footer is written). Whatever was already written
+    # to the IO stays there; discarding it is up to the caller.
     def abort
-      return if @closed
       @closed = true
-      return unless @temp_path
-      @io.close unless @io.closed?
-      File.unlink(@temp_path) if File.exist?(@temp_path)
     end
 
     private
 
+    # Pathname responds to #write too (writing a whole file by path), so it is rejected explicitly
+    def check_io!(io)
+      return io if io.respond_to?(:write) && !io.is_a?(String) && !(defined?(Pathname) && io.is_a?(Pathname))
+      raise ArgumentError, "Herringbone::Writer expects an IO that responds to #write " \
+        "(e.g. File.open(path, \"wb\") or StringIO.new), got #{io.class}"
+    end
+
     # Levels are kept as binary Strings (one byte per entry). Values are an Array for numeric and
     # boolean columns, and a compact ByteValues for BYTE_ARRAY / FIXED_LEN_BYTE_ARRAY columns.
     ColumnBuffer = Struct.new(:defs, :reps, :values)
@@ -220,24 +237,7 @@ module Herringbone
       end
     end
 
-    def open_target(target)
-      if target.respond_to?(:write)
-        @io = target
-      else
-        @path = target.to_s
-        dir = File.dirname(@path)
-        @temp_path = File.join(dir, ".#{File.basename(@path)}.#{Process.pid}.#{rand(1 << 32).to_s(36)}.tmp")
-        @io = File.open(@temp_path, "wb")
-      end
-    end
-
-    def finish_target
-      return unless @temp_path
-      @io.close
-      File.rename(@temp_path, @path)
-    end
-
-    # Flushes when the buffered values reach row_group_bytes (or row_group_size rows). The bytes per
+    # Flushes when the buffered values reach row_group_bytes (or row_group_rows rows). The bytes per
     # row are estimated from the buffered values after the first rows, then refreshed per row group.
     def check_row_group_size
       @bytes_per_row ||= estimate_bytes_per_row
@@ -253,7 +253,7 @@ module Herringbone
 
     def row_limit_for(bytes_per_row)
       by_bytes = [@row_group_bytes / [bytes_per_row, 1].max, ESTIMATE_AFTER_ROWS].max
-      @row_group_size ? [@row_group_size, by_bytes].min : by_bytes
+      @row_group_rows ? [@row_group_rows, by_bytes].min : by_bytes
     end
 
     def estimate_bytes_per_row
@@ -430,6 +430,9 @@ module Herringbone
       else
         values.sum(&:bytesize) + 4 * values.size
       end
+      order = sort_key(col)
+      pages = []
+      first_row = 0
       page_ranges(buf, value_bytes).each do |from, to|
         n = to - from
         defs = buf.defs[from, n]
@@ -437,6 +440,12 @@ module Herringbone
         non_null = max_def.zero? ? n : defs.count(max_def)
         page_values = (indices || values)[value_index, non_null]
         value_index += non_null
+        page_offset = @pos
+        range = nil
+        if order && non_null.positive?
+          in_page = dict_values ? page_values.uniq.map { |i| dict_values[i] } : page_values
+          range = value_range(col, in_page, order)
+        end
         encoded = encode_values(page_values, value_encoding, type, col.type_length, dict_values&.size)
         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
@@ -446,6 +455,8 @@ module Herringbone
           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)
         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)
       end
 
       encodings = [E::RLE]
@@ -461,9 +472,125 @@ module Herringbone
         total_compressed_size: @pos - chunk_start,
         data_page_offset: data_offset,
         dictionary_page_offset: dictionary_offset,
-        statistics: @statistics ? statistics_for(col, buf.defs, values || indices.uniq.map { |i| dict_values[i] }) : nil
+        statistics: statistics_for(col, buf.defs, values || indices.uniq.map { |i| dict_values[i] }, order)
       )
-      Format::ColumnChunk.new(file_offset: chunk_start, meta_data: meta)
+      chunk = Format::ColumnChunk.new(file_offset: chunk_start, meta_data: meta)
+      if (bloom = @bloom_filters[path])
+        @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)]
+      chunk
+    end
+
+    PageInfo = Struct.new(:offset, :size, :first_row, :nulls, :non_null, :range)
+
+    def offset_index_for(pages)
+      Format::OffsetIndex.new(page_locations: pages.map do |p|
+        Format::PageLocation.new(offset: p.offset, compressed_page_size: p.size, first_row_index: p.first_row)
+      end)
+    end
+
+    # nil when the column has no defined sort order, or a page's values have no min/max (all NaN)
+    def column_index_for(col, pages, order)
+      return nil unless order
+      return nil if pages.any? { |p| p.non_null.positive? && p.range.nil? }
+      Format::ColumnIndex.new(
+        null_pages: pages.map { |p| p.non_null.zero? },
+        min_values: pages.map { |p| p.range ? truncate_min(stat_bytes(col, p.range[0])) : "".b },
+        max_values: pages.map { |p| p.range ? truncate_max(stat_bytes(col, p.range[1])) : "".b },
+        boundary_order: boundary_order(pages.filter_map(&:range), order),
+        null_counts: pages.map(&:nulls)
+      )
+    end
+
+    def boundary_order(ranges, order)
+      return Format::BoundaryOrder::ASCENDING if ranges.size < 2
+      cmp = ->(a, b) { order.equal?(IDENTITY) ? a <=> b : order.call(a) <=> order.call(b) }
+      pairs = ranges.each_cons(2)
+      if pairs.all? { |(a_min, a_max), (b_min, b_max)| cmp.(a_min, b_min) <= 0 && cmp.(a_max, b_max) <= 0 }
+        Format::BoundaryOrder::ASCENDING
+      elsif pairs.all? { |(a_min, a_max), (b_min, b_max)| cmp.(a_min, b_min) >= 0 && cmp.(a_max, b_max) >= 0 }
+        Format::BoundaryOrder::DESCENDING
+      else
+        Format::BoundaryOrder::UNORDERED
+      end
+    end
+
+    # Page indexes go after the last row group: all column indexes, then all offset indexes
+    def write_page_indexes
+      @page_indexes.each do |chunk, column_index, _|
+        next unless column_index
+        bytes = 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|
+        bytes = offset_index.encode
+        chunk.offset_index_offset = @pos
+        chunk.offset_index_length = bytes.bytesize
+        write_raw(bytes)
+      end
+      @page_indexes.clear
+    end
+
+    BLOOM_FILTER_OPTIONS = %i[ndv fpp max_bytes].freeze
+
+    # { dotted_path => { ndv:, fpp:, max_bytes: } } from the bloom_filters: option
+    def bloom_filter_config(option)
+      return {} if option.nil? || option == false
+      columns = @schema.columns.to_h { |c| [c.dotted_path, c] }
+      if option == true
+        return columns.select { |_, c| BloomFilter::TYPES.include?(c.type) }.transform_values { bloom_filter_settings({}) }
+      end
+      option = Array(option).to_h { |path| [path, true] } unless option.is_a?(Hash)
+      option.each_with_object({}) do |(path, settings), config|
+        path = path.is_a?(Array) ? path.join(".") : path.to_s
+        col = columns[path] or raise ArgumentError, "bloom_filters: no such column #{path}"
+        next if settings.nil? || settings == false
+        unless BloomFilter::TYPES.include?(col.type)
+          raise ArgumentError, "bloom_filters: not supported for #{T::NAMES[col.type]} column #{path}"
+        end
+        settings = {} if settings == true
+        raise ArgumentError, "bloom_filters: expected true or a Hash for #{path}, got #{settings.inspect}" unless settings.is_a?(Hash)
+        config[path] = bloom_filter_settings(settings.transform_keys(&:to_sym), path)
+      end
+    end
+
+    def bloom_filter_settings(settings, path = nil)
+      unknown = settings.keys - BLOOM_FILTER_OPTIONS
+      raise ArgumentError, "bloom_filters: unknown option #{unknown.join(", ")} for #{path}" unless unknown.empty?
+      ndv = settings[:ndv] && Integer(settings[:ndv])
+      raise ArgumentError, "bloom_filters: ndv must be positive for #{path}" if ndv && !ndv.positive?
+      fpp = Float(settings.fetch(:fpp, BloomFilter::DEFAULT_FPP))
+      raise ArgumentError, "bloom_filters: fpp must be between 0 and 1 for #{path}" unless fpp > 0 && fpp < 1
+      { ndv: ndv, fpp: fpp, max_bytes: Integer(settings.fetch(:max_bytes, BloomFilter::DEFAULT_MAX_BYTES)) }
+    end
+
+    # A filter holding the chunk's values: +dict_values+ (already distinct) for dictionary-encoded
+    # chunks, +values+ otherwise. Each distinct value is hashed once, and the filter is sized from
+    # the configured ndv or from the number of distinct values.
+    def build_bloom_filter(col, settings, dict_values, values)
+      hashes = if dict_values
+        BloomFilter.hash_physical_all(dict_values, col.type)
+      else
+        BloomFilter.hash_physical_all(values, col.type, distinct: true)
+      end
+      size = BloomFilter.optimal_num_bytes(settings[:ndv] || hashes.size, settings[:fpp], max_bytes: settings[:max_bytes])
+      filter = BloomFilter.new(size, column: col)
+      filter.insert_hashes(hashes)
+      filter
+    end
+
+    # Bloom filters go right after the row group's column chunks, in column order
+    def write_bloom_filters
+      @pending_bloom_filters.each do |meta, filter|
+        bytes = filter.encode
+        meta.bloom_filter_offset = @pos
+        meta.bloom_filter_length = bytes.bytesize
+        write_raw(bytes)
+      end
+      @pending_bloom_filters.clear
     end
 
     def use_dictionary?(col)
@@ -497,14 +624,15 @@ module Herringbone
       [dict, values.map(&index)]
     end
 
-    # Splits a column buffer into pages of roughly @page_size bytes. Repeated columns
+    # Splits a column buffer into pages of roughly @page_bytes bytes. Repeated columns
     # are only cut where a new row starts.
     def page_ranges(buf, value_bytes)
       n = buf.defs.size
       bytes = n + value_bytes
-      pages = (bytes + @page_size - 1) / @page_size
-      return [[0, n]] if pages <= 1
-      per = (n + pages - 1) / pages
+      pages = (bytes + @page_bytes - 1) / @page_bytes
+      per = pages <= 1 ? n : (n + pages - 1) / pages
+      per = @page_rows if per > @page_rows
+      return [[0, n]] if per >= n
       reps = buf.reps
       ranges = []
       start = 0
@@ -589,46 +717,86 @@ module Herringbone
       write_page(header, "".b, rep_bytes + def_bytes + compressed)
     end
 
-    def statistics_for(col, defs, values)
+    STAT_TRUNCATE_BYTES = 64
+    IDENTITY = ->(v) { v }
+
+    def statistics_for(col, defs, values, order)
       max_def = col.max_definition_level
       nulls = max_def.zero? ? 0 : defs.size - defs.count(max_def)
       stats = Format::Statistics.new(null_count: nulls)
-      min, max = min_max(col, values)
-      if min
-        stats.min_value = min
-        stats.max_value = max
-        stats.is_min_value_exact = true
-        stats.is_max_value_exact = true
+      range = order && value_range(col, values, order)
+      if range
+        min = stat_bytes(col, range[0])
+        max = stat_bytes(col, range[1])
+        stats.min_value = truncate_min(min)
+        stats.max_value = truncate_max(max)
+        stats.is_min_value_exact = stats.min_value.bytesize == min.bytesize
+        stats.is_max_value_exact = stats.max_value.bytesize == max.bytesize
       end
       stats
     end
 
-    # Min/max for types whose Parquet sort order matches Ruby's comparison of the physical values
-    def min_max(col, values)
-      return nil if values.empty?
+    # A key giving the Parquet sort order of the column's physical values (IDENTITY when Ruby's own
+    # comparison already matches), or nil when the order is undefined (INT96)
+    def sort_key(col)
       kind, _, signed = Types.logical_of(col.node)
       case col.type
-      when T::BOOLEAN
-        [values.include?(false) ? "\x00".b : "\x01".b, values.include?(true) ? "\x01".b : "\x00".b]
+      when T::BOOLEAN then ->(v) { v ? 1 : 0 }
       when T::INT32, T::INT64
-        return nil if kind == :integer && !signed
-        fmt = col.type == T::INT32 ? "l<" : "q<"
-        min, max = values.minmax
-        [[min].pack(fmt), [max].pack(fmt)]
-      when T::FLOAT, T::DOUBLE
-        finite = values.reject(&:nan?)
-        return nil if finite.empty?
-        min, max = finite.minmax
+        return IDENTITY unless kind == :integer && !signed
+        mask = col.type == T::INT32 ? 0xFFFF_FFFF : 0xFFFF_FFFF_FFFF_FFFF
+        ->(v) { v & mask }
+      when T::FLOAT, T::DOUBLE then IDENTITY
+      when T::BYTE_ARRAY, T::FIXED_LEN_BYTE_ARRAY
+        case kind
+        when :decimal then Types.method(:be_to_int)
+        when :float16 then ->(v) { Types.half_to_float(v.unpack1("v")) }
+        else IDENTITY # unsigned lexicographic, which is how Ruby compares binary Strings
+        end
+      end
+    end
+
+    # [min, max] of +values+ in column order, ignoring NaNs; nil if there is nothing to compare
+    def value_range(col, values, order)
+      floats = col.type == T::FLOAT || col.type == T::DOUBLE
+      if floats
+        values = values.reject(&:nan?)
+      elsif Types.logical_of(col.node).first == :float16
+        values = values.reject { |v| order.call(v).nan? }
+      end
+      return nil if values.empty?
+      min, max = order.equal?(IDENTITY) ? values.minmax : values.minmax_by(&order)
+      if floats
         min = -0.0 if min.zero?
         max = 0.0 if max.zero?
-        fmt = col.type == T::FLOAT ? "e" : "E"
-        [[min].pack(fmt), [max].pack(fmt)]
-      when T::BYTE_ARRAY
-        return nil if kind == :decimal
-        min, max = values.minmax
-        return nil if min.bytesize > 1024 || max.bytesize > 1024
-        [min, max]
       end
+      [min, max]
+    end
+
+    def stat_bytes(col, value)
+      case col.type
+      when T::BOOLEAN then value ? "\x01".b : "\x00".b
+      when T::INT32 then [value].pack("l<")
+      when T::INT64 then [value].pack("q<")
+      when T::FLOAT then [value].pack("e")
+      when T::DOUBLE then [value].pack("E")
+      else value.b
+      end
+    end
+
+    # Long byte-array bounds are truncated: a prefix is still a lower bound for the minimum, and
+    # a prefix with its last byte incremented is an upper bound for the maximum
+    def truncate_min(bytes)
+      bytes.bytesize > STAT_TRUNCATE_BYTES ? bytes.byteslice(0, STAT_TRUNCATE_BYTES) : bytes
+    end
+
+    def truncate_max(bytes)
+      return bytes if bytes.bytesize <= STAT_TRUNCATE_BYTES
+      prefix = bytes.byteslice(0, STAT_TRUNCATE_BYTES).bytes
+      prefix.pop while prefix.last == 0xFF
+      return bytes if prefix.empty?
+      prefix[-1] += 1
+      prefix.pack("C*")
     end
   end
 end