herringbone 0.1.0

data/lib/herringbone/compression.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+require "zlib"
+require "stringio"
+require "zstd-ruby"
+require "brotli"
+
+module Herringbone
+  # Dispatches page (de)compression by Parquet codec id. Snappy and LZ4 are pure Ruby,
+  # GZIP uses zlib, ZSTD and Brotli use the zstd-ruby and brotli gems.
+  module Compression
+    module_function
+
+    CODECS_BY_NAME = {
+      none: Format::Codec::UNCOMPRESSED, uncompressed: Format::Codec::UNCOMPRESSED,
+      snappy: Format::Codec::SNAPPY, gzip: Format::Codec::GZIP, brotli: Format::Codec::BROTLI,
+      lz4: Format::Codec::LZ4_RAW, lz4_raw: Format::Codec::LZ4_RAW, lz4_hadoop: Format::Codec::LZ4,
+      zstd: Format::Codec::ZSTD
+    }.freeze
+
+    def codec_id(name)
+      return name if name.is_a?(Integer)
+      CODECS_BY_NAME.fetch(name.to_s.downcase.to_sym) { raise ArgumentError, "Unknown compression codec #{name.inspect}" }
+    end
+
+    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::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)
+      when Format::Codec::LZ4 then Codecs::LZ4.decompress_hadoop(data, uncompressed_size)
+      when Format::Codec::ZSTD then ::Zstd.decompress(data)
+      when Format::Codec::BROTLI then ::Brotli.inflate(data)
+      else
+        raise UnsupportedError, "Unsupported compression codec #{Format::Codec::NAMES[codec] || codec}"
+      end
+      out = out.b
+      if out.bytesize != uncompressed_size
+        raise DecodeError, "Decompressed #{out.bytesize} bytes, expected #{uncompressed_size}"
+      end
+      out
+    end
+
+    def compress(codec, data)
+      case codec
+      when Format::Codec::UNCOMPRESSED then data
+      when Format::Codec::SNAPPY then Codecs::Snappy.compress(data)
+      when Format::Codec::GZIP then Zlib.gzip(data)
+      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 ::Zstd.compress(data)
+      when Format::Codec::BROTLI then ::Brotli.deflate(data)
+      else
+        raise UnsupportedError, "Unsupported compression codec #{Format::Codec::NAMES[codec] || codec}"
+      end.b
+    end
+
+    # Handles files whose gzip data consists of several concatenated members
+    def gunzip(data)
+      io = StringIO.new(data)
+      out = String.new(encoding: Encoding::BINARY)
+      loop do
+        gz = Zlib::GzipReader.new(io)
+        out << gz.read
+        unused = gz.unused
+        gz.finish
+        break if unused.nil? || unused.empty?
+        io.pos -= unused.bytesize
+      end
+      out
+    end
+  end
+end

data/lib/herringbone/encodings/delta.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module Herringbone
+  module Encodings
+    # DELTA_BINARY_PACKED, DELTA_LENGTH_BYTE_ARRAY and DELTA_BYTE_ARRAY
+    module Delta
+      module_function
+
+      BLOCK_SIZE = 128
+      MINIBLOCKS = 4
+
+      def zigzag_decode(n) = (n >> 1) ^ -(n & 1)
+      def zigzag_encode(n) = n.negative? ? ((-n) << 1) - 1 : n << 1
+
+      # Wraps an integer into the signed range of +bits+ bits
+      def wrap(v, bits)
+        half = 1 << (bits - 1)
+        ((v + half) & ((1 << bits) - 1)) - half
+      end
+
+      # Decodes DELTA_BINARY_PACKED integers. +bits+ is 32 or 64 (for wraparound).
+      # Returns [values, new_pos]. If +count+ is nil, the total from the header is used.
+      def decode_binary_packed(data, pos, bits = 64, count = nil)
+        block_size, pos = RLE.read_uleb(data, pos)
+        miniblocks, pos = RLE.read_uleb(data, pos)
+        total, pos = RLE.read_uleb(data, pos)
+        first, pos = RLE.read_uleb(data, pos)
+        raise DecodeError, "Invalid DELTA_BINARY_PACKED header" if miniblocks.zero? || block_size % miniblocks != 0
+        per_mini = block_size / miniblocks
+        raise DecodeError, "Invalid miniblock size #{per_mini}" if per_mini % 8 != 0
+        total = count if count && count < total
+        values = []
+        return [values, pos] if total.zero?
+        last = zigzag_decode(first)
+        values << last
+        half = 1 << (bits - 1)
+        mask = (1 << bits) - 1
+        while values.size < total
+          min_delta, pos = RLE.read_uleb(data, pos)
+          min_delta = zigzag_decode(min_delta)
+          widths = data.byteslice(pos, miniblocks).unpack("C*")
+          pos += miniblocks
+          widths.each do |w|
+            break if values.size >= total
+            raise DecodeError, "Invalid delta bit width #{w}" if w > bits
+            deltas = RLE.unpack_bits(data, pos, per_mini, w)
+            pos += per_mini * w / 8
+            take = total - values.size
+            deltas = deltas.first(take) if take < per_mini
+            deltas.each do |d|
+              last = ((last + min_delta + d + half) & mask) - half
+              values << last
+            end
+          end
+        end
+        [values, pos]
+      end
+
+      def encode_binary_packed(values, bits = 64)
+        out = String.new(encoding: Encoding::BINARY)
+        per_mini = BLOCK_SIZE / MINIBLOCKS
+        RLE.write_uleb(out, BLOCK_SIZE)
+        RLE.write_uleb(out, MINIBLOCKS)
+        RLE.write_uleb(out, values.size)
+        RLE.write_uleb(out, zigzag_encode(values.empty? ? 0 : values[0]))
+        return out if values.size <= 1
+
+        mask = (1 << bits) - 1
+        deltas = Array.new(values.size - 1) { |i| wrap(values[i + 1] - values[i], bits) }
+        deltas.each_slice(BLOCK_SIZE) do |block|
+          min = block.min
+          RLE.write_uleb(out, zigzag_encode(min))
+          adjusted = block.map { |d| (d - min) & mask }
+          minis = adjusted.each_slice(per_mini).to_a
+          widths = Array.new(MINIBLOCKS) { |m| minis[m] ? minis[m].max.bit_length : 0 }
+          out << widths.pack("C*")
+          minis.each_with_index do |mini, m|
+            mini += Array.new(per_mini - mini.size, 0) if mini.size < per_mini
+            out << RLE.pack_bits(mini, widths[m])
+          end
+        end
+        out
+      end
+
+      def decode_length_byte_array(data, pos, count)
+        lengths, pos = decode_binary_packed(data, pos, 32, count)
+        raise DecodeError, "DELTA_LENGTH_BYTE_ARRAY has #{lengths.size} lengths, need #{count}" if lengths.size < count
+        out = Array.new(count)
+        size = data.bytesize
+        lengths.each_with_index do |len, i|
+          raise DecodeError, "DELTA_LENGTH_BYTE_ARRAY value overruns page" if len.negative? || pos + len > size
+          out[i] = data.byteslice(pos, len)
+          pos += len
+        end
+        [out, pos]
+      end
+
+      def encode_length_byte_array(values)
+        encode_binary_packed(values.map(&:bytesize), 32) << values.pack("a*" * values.size)
+      end
+
+      def decode_byte_array(data, pos, count)
+        prefixes, pos = decode_binary_packed(data, pos, 32, count)
+        suffixes, pos = decode_length_byte_array(data, pos, count)
+        prev = "".b
+        out = Array.new(count) do |i|
+          prefix = prefixes[i]
+          raise DecodeError, "DELTA_BYTE_ARRAY prefix longer than previous value" if prefix > prev.bytesize
+          prev = prefix.zero? ? suffixes[i] : prev.byteslice(0, prefix) + suffixes[i]
+        end
+        [out, pos]
+      end
+
+      def encode_byte_array(values)
+        prev = "".b
+        prefixes = []
+        suffixes = []
+        values.each do |v|
+          max = [prev.bytesize, v.bytesize].min
+          n = 0
+          n += 1 while n < max && prev.getbyte(n) == v.getbyte(n)
+          prefixes << n
+          suffixes << v.byteslice(n, v.bytesize - n)
+          prev = v
+        end
+        encode_binary_packed(prefixes, 32) << encode_length_byte_array(suffixes)
+      end
+    end
+
+    # BYTE_STREAM_SPLIT: byte k of every value is stored in stream k.
+    module ByteStreamSplit
+      module_function
+
+      # Returns the value bytes re-interleaved into PLAIN layout
+      def decode(data, pos, count, width)
+        nbytes = count * width
+        raise DecodeError, "Truncated BYTE_STREAM_SPLIT data" if pos + nbytes > data.bytesize
+        streams = Array.new(width) { |k| data.byteslice(pos + k * count, count).unpack("C*") }
+        [streams[0].zip(*streams[1..]).flatten.pack("C*"), pos + nbytes]
+      end
+
+      def encode(plain, width)
+        count = plain.bytesize / width
+        bytes = plain.unpack("C*")
+        out = String.new(capacity: plain.bytesize, encoding: Encoding::BINARY)
+        width.times do |k|
+          out << Array.new(count) { |i| bytes[i * width + k] }.pack("C*")
+        end
+        out
+      end
+    end
+  end
+end

data/lib/herringbone/encodings/plain.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+module Herringbone
+  module Encodings
+    # PLAIN encoding for all physical types. Values are returned in their physical
+    # Ruby form (Integer, Float, true/false, binary String); logical conversion happens elsewhere.
+    module Plain
+      module_function
+
+      FORMATS = {
+        Format::Type::INT32 => ["l<", 4],
+        Format::Type::INT64 => ["q<", 8],
+        Format::Type::FLOAT => ["e", 4],
+        Format::Type::DOUBLE => ["E", 8]
+      }.freeze
+
+      # Decodes +count+ values of +type+ from +data+ starting at +pos+.
+      # Returns [values, new_pos].
+      def decode(data, pos, count, type, type_length = nil)
+        case type
+        when Format::Type::BOOLEAN
+          nbytes = (count + 7) / 8
+          bits = data.byteslice(pos, nbytes).unpack1("b*")
+          raise DecodeError, "Truncated BOOLEAN data" if bits.bytesize < count
+          [Array.new(count) { |i| bits.getbyte(i) == 49 }, pos + nbytes]
+        when Format::Type::INT32, Format::Type::INT64, Format::Type::FLOAT, Format::Type::DOUBLE
+          fmt, width = FORMATS[type]
+          nbytes = count * width
+          raise DecodeError, "Truncated PLAIN data" if pos + nbytes > data.bytesize
+          [data.byteslice(pos, nbytes).unpack("#{fmt}#{count}"), pos + nbytes]
+        when Format::Type::INT96
+          nbytes = count * 12
+          raise DecodeError, "Truncated INT96 data" if pos + nbytes > data.bytesize
+          vals = data.byteslice(pos, nbytes).unpack("Q<L<" * count).each_slice(2).map { |nanos, day| [nanos, day] }
+          [vals, pos + nbytes]
+        when Format::Type::BYTE_ARRAY
+          decode_byte_arrays(data, pos, count)
+        when Format::Type::FIXED_LEN_BYTE_ARRAY
+          nbytes = count * type_length
+          raise DecodeError, "Truncated FIXED_LEN_BYTE_ARRAY data" if pos + nbytes > data.bytesize
+          [Array.new(count) { |i| data.byteslice(pos + i * type_length, type_length) }, pos + nbytes]
+        else
+          raise DecodeError, "Unknown physical type #{type}"
+        end
+      end
+
+      def decode_byte_arrays(data, pos, count)
+        out = Array.new(count)
+        size = data.bytesize
+        i = 0
+        while i < count
+          raise DecodeError, "Truncated BYTE_ARRAY data" if pos + 4 > size
+          len = data.getbyte(pos) | (data.getbyte(pos + 1) << 8) | (data.getbyte(pos + 2) << 16) | (data.getbyte(pos + 3) << 24)
+          pos += 4
+          raise DecodeError, "BYTE_ARRAY value overruns page" if pos + len > size
+          out[i] = data.byteslice(pos, len)
+          pos += len
+          i += 1
+        end
+        [out, pos]
+      end
+
+      def encode(values, type, type_length = nil)
+        case type
+        when Format::Type::BOOLEAN
+          [values.map { |v| v ? "1" : "0" }.join].pack("b*")
+        when Format::Type::INT32, Format::Type::INT64, Format::Type::FLOAT, Format::Type::DOUBLE
+          values.pack("#{FORMATS[type][0]}*")
+        when Format::Type::INT96
+          values.flat_map { |nanos, day| [nanos, day] }.pack("Q<L<" * values.size)
+        when Format::Type::BYTE_ARRAY
+          # pack("a*") copies raw bytes whatever the string's encoding, without an intermediate copy
+          args = []
+          values.each { |v| args << v.bytesize << v }
+          args.pack("Va*" * values.size)
+        when Format::Type::FIXED_LEN_BYTE_ARRAY
+          values.each do |v|
+            raise EncodeError, "Expected #{type_length} bytes, got #{v.bytesize}" if v.bytesize != type_length
+          end
+          values.pack("a*" * values.size)
+        else
+          raise EncodeError, "Unknown physical type #{type}"
+        end
+      end
+    end
+  end
+end

data/lib/herringbone/encodings/rle.rb

@@ -0,0 +1,203 @@
+# frozen_string_literal: true
+
+module Herringbone
+  module Encodings
+    # Bit packing (LSB-first, as used by Parquet) and the RLE / bit-packed hybrid encoding
+    # used for repetition/definition levels, dictionary indices and RLE booleans.
+    module RLE
+      module_function
+
+      # Unpacks +count+ values of +width+ bits each, starting at byte +offset+ of +data+.
+      # Missing trailing bytes are treated as zeroes.
+      def unpack_bits(data, offset, count, width)
+        return Array.new(count, 0) if width.zero?
+        nbytes = (count * width + 7) / 8
+        chunk = data.byteslice(offset, nbytes) || "".b
+        if width == 8
+          vals = chunk.unpack("C*")
+          vals.fill(0, vals.size, count - vals.size) if vals.size < count
+          return vals
+        end
+        return unpack_wide(chunk, count, width) if width > 32
+
+        # Pad to a whole number of 32-bit words plus one spare word, so reads never go out of range
+        pad = (-chunk.bytesize % 4) + 4
+        words = (pad.zero? ? chunk : chunk + ("\0" * pad)).unpack("V*")
+        mask = (1 << width) - 1
+        out = Array.new(count)
+        bitpos = 0
+        i = 0
+        while i < count
+          wi = bitpos >> 5
+          off = bitpos & 31
+          v = words[wi] >> off
+          spill = off + width - 32
+          v |= (words[wi + 1] & ((1 << spill) - 1)) << (32 - off) if spill.positive?
+          out[i] = v & mask
+          bitpos += width
+          i += 1
+        end
+        out
+      end
+
+      # Slow path for widths above 32 bits (used by DELTA_BINARY_PACKED with 64-bit values)
+      def unpack_wide(chunk, count, width)
+        mask = (1 << width) - 1
+        out = Array.new(count)
+        group_bytes = width # 8 values * width bits = width bytes
+        ngroups = (count + 7) / 8
+        i = 0
+        ngroups.times do |g|
+          bytes = chunk.byteslice(g * group_bytes, group_bytes) || "".b
+          bytes += "\0" * (group_bytes - bytes.bytesize) if bytes.bytesize < group_bytes
+          big = bytes.reverse.unpack1("H*").to_i(16)
+          8.times do |j|
+            break if i >= count
+            out[i] = (big >> (j * width)) & mask
+            i += 1
+          end
+        end
+        out
+      end
+
+      # Packs +values+ with +width+ bits each; the value count is padded up to a multiple of 8.
+      def pack_bits(values, width)
+        return "".b if width.zero? || values.empty?
+        n = values.size
+        padded = (n + 7) & ~7
+        if width == 8
+          s = values.pack("C*")
+          s << ("\0" * (padded - n)) if padded > n
+          return s
+        end
+        out = String.new(capacity: padded * width / 8, encoding: Encoding::BINARY)
+        acc = 0
+        nbits = 0
+        i = 0
+        while i < padded
+          acc |= (values[i] || 0) << nbits
+          nbits += width
+          while nbits >= 32
+            out << [acc & 0xFFFFFFFF].pack("V")
+            acc >>= 32
+            nbits -= 32
+          end
+          i += 1
+        end
+        # Remaining bits: padded*width is a multiple of 8, so nbits is a whole number of bytes
+        while nbits.positive?
+          out << (acc & 0xFF)
+          acc >>= 8
+          nbits -= 8
+        end
+        out
+      end
+
+      def read_uleb(data, pos)
+        result = 0
+        shift = 0
+        while true
+          b = data.getbyte(pos)
+          raise DecodeError, "Truncated varint" unless b
+          pos += 1
+          result |= (b & 0x7F) << shift
+          return [result, pos] if b < 0x80
+          shift += 7
+        end
+      end
+
+      def write_uleb(out, n)
+        while n >= 0x80
+          out << ((n & 0x7F) | 0x80)
+          n >>= 7
+        end
+        out << n
+      end
+
+      # Decodes the RLE/bit-packed hybrid from +data+ between +pos+ and +limit+,
+      # returning exactly +count+ values (missing values are an error).
+      def decode_hybrid(data, pos, limit, width, count)
+        out = []
+        value_bytes = (width + 7) / 8
+        while out.size < count
+          raise DecodeError, "RLE data exhausted (#{out.size}/#{count} values)" if pos >= limit
+          header, pos = read_uleb(data, pos)
+          if header & 1 == 1
+            groups = header >> 1
+            n = groups * 8
+            vals = unpack_bits(data, pos, n, width)
+            pos += groups * width
+            remaining = count - out.size
+            vals = vals.first(remaining) if vals.size > remaining
+            out.concat(vals)
+          else
+            run = header >> 1
+            v = 0
+            value_bytes.times { |k| v |= data.getbyte(pos + k).to_i << (8 * k) }
+            pos += value_bytes
+            remaining = count - out.size
+            run = remaining if run > remaining
+            out.fill(v, out.size, run)
+          end
+        end
+        out
+      end
+
+      # Encodes +values+ with the RLE/bit-packed hybrid. Repeated runs of 8+ equal values
+      # become RLE runs, everything else goes into bit-packed groups of 8.
+      def encode_hybrid(values, width)
+        out = String.new(encoding: Encoding::BINARY)
+        n = values.size
+        value_bytes = (width + 7) / 8
+        return out if n.zero?
+        min, max = values.minmax
+        if min == max
+          # A single run, the common case for levels of columns without nulls
+          write_uleb(out, n << 1)
+          value_bytes.times { |k| out << ((min >> (8 * k)) & 0xFF) }
+          return out
+        end
+        literal_start = nil
+        i = 0
+        while i < n
+          v = values[i]
+          # Measure how long the run of v starting at i is
+          j = i + 1
+          j += 1 while j < n && values[j] == v
+          run = j - i
+          if run >= 8 && (literal_start.nil? || (i - literal_start) % 8 == 0)
+            flush_literals(out, values, literal_start, i, width) if literal_start
+            literal_start = nil
+            write_uleb(out, run << 1)
+            value_bytes.times { |k| out << ((v >> (8 * k)) & 0xFF) }
+            i = j
+          else
+            literal_start ||= i
+            # Consume values so that the literal count stays aligned to groups of 8 where possible
+            i += 1
+          end
+        end
+        flush_literals(out, values, literal_start, n, width) if literal_start
+        out
+      end
+
+      def flush_literals(out, values, from, to, width)
+        count = to - from
+        groups = (count + 7) / 8
+        write_uleb(out, (groups << 1) | 1)
+        out << pack_bits(values[from, count], width)
+      end
+
+      def bit_width(max_value)
+        max_value.to_i.bit_length
+      end
+
+      # Legacy BIT_PACKED level encoding (deprecated): MSB-first bit order, no header.
+      def decode_legacy_bit_packed(data, pos, width, count)
+        nbytes = (count * width + 7) / 8
+        bits = data.byteslice(pos, nbytes).unpack1("B*")
+        Array.new(count) { |i| bits[i * width, width].to_i(2) }
+      end
+    end
+  end
+end