herringbone 0.1.0

data/lib/herringbone/types.rb

@@ -0,0 +1,454 @@
+# frozen_string_literal: true
+
+require "date"
+require "time"
+require "json"
+require "bigdecimal"
+
+module Herringbone
+  # Conversion between physical Parquet values and Ruby objects, driven by the
+  # logical type (or legacy converted type) of a column.
+  #
+  #   STRING/ENUM/JSON      <-> String (UTF-8)
+  #   BYTE_ARRAY/FLBA/BSON  <-> String (binary)
+  #   INTEGER (signed/uns.) <-> Integer
+  #   DATE                  <-> Date
+  #   TIMESTAMP, INT96      <-> Time (UTC)
+  #   TIME                  <-> Integer in the column's unit since midnight
+  #   DECIMAL               <-> BigDecimal
+  #   UUID                  <-> String "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
+  #   FLOAT16               <-> Float
+  module Types
+    module_function
+
+    T = Format::Type
+    C = Format::ConvertedType
+
+    EPOCH_JD = Date.new(1970, 1, 1).jd
+    JULIAN_EPOCH_DAY = 2_440_588 # Julian day number of 1970-01-01, used by INT96
+    NANOS_PER_DAY = 86_400 * 1_000_000_000
+
+    UNIT_DIVISORS = { millis: 1_000, micros: 1_000_000, nanos: 1_000_000_000 }.freeze
+    UNIT_NAMES = { millis: :millisecond, micros: :microsecond, nanos: :nanosecond }.freeze
+
+
+    def lt(**kw) = Format::LogicalType.new(**kw)
+
+    def time_unit(unit)
+      case unit.to_sym
+      when :millis, :ms then Format::TimeUnit.millis
+      when :micros, :us then Format::TimeUnit.micros
+      when :nanos, :ns then Format::TimeUnit.nanos
+      else raise ArgumentError, "Unknown time unit #{unit.inspect}"
+      end
+    end
+
+    def int_type(bits, signed)
+      physical = bits == 64 ? T::INT64 : T::INT32
+      converted = signed ? C.const_get("INT_#{bits}") : C.const_get("UINT_#{bits}")
+      { type: physical, converted_type: converted,
+        logical_type: lt(integer: Format::IntType.new(bit_width: bits, is_signed: signed)) }
+    end
+
+    # Physical attributes (type, logical type etc.) for a DSL type name
+    def physical_attributes(type, **opts)
+      case type.to_sym
+      when :boolean, :bool then { type: T::BOOLEAN }
+      when :int8 then int_type(8, true)
+      when :int16 then int_type(16, true)
+      when :int32 then { type: T::INT32 }
+      when :int64 then { type: T::INT64 }
+      when :uint8 then int_type(8, false)
+      when :uint16 then int_type(16, false)
+      when :uint32 then int_type(32, false)
+      when :uint64 then int_type(64, false)
+      when :float then { type: T::FLOAT }
+      when :double then { type: T::DOUBLE }
+      when :float16 then { type: T::FIXED_LEN_BYTE_ARRAY, type_length: 2, logical_type: lt(float16: Format::Float16Type.new) }
+      when :string, :utf8 then { type: T::BYTE_ARRAY, converted_type: C::UTF8, logical_type: lt(string: Format::StringType.new) }
+      when :binary, :bytes then { type: T::BYTE_ARRAY }
+      when :json then { type: T::BYTE_ARRAY, converted_type: C::JSON, logical_type: lt(json: Format::JsonType.new) }
+      when :bson then { type: T::BYTE_ARRAY, converted_type: C::BSON, logical_type: lt(bson: Format::BsonType.new) }
+      when :enum
+        if opts[:parquet_enum]
+          { type: T::BYTE_ARRAY, converted_type: C::ENUM, logical_type: lt(enum: Format::EnumType.new) }
+        else
+          physical_attributes(:string)
+        end
+      when :uuid then { type: T::FIXED_LEN_BYTE_ARRAY, type_length: 16, logical_type: lt(uuid: Format::UUIDType.new) }
+      when :date then { type: T::INT32, converted_type: C::DATE, logical_type: lt(date: Format::DateType.new) }
+      when :int96 then { type: T::INT96 }
+      when :fixed
+        { type: T::FIXED_LEN_BYTE_ARRAY, type_length: Integer(opts.fetch(:length)) }
+      when :time
+        unit = time_unit(opts.fetch(:unit, :micros))
+        converted = { millis: C::TIME_MILLIS, micros: C::TIME_MICROS }[unit.to_sym]
+        { type: unit.millis ? T::INT32 : T::INT64, converted_type: opts.fetch(:utc, true) ? converted : nil,
+          logical_type: lt(time: Format::TimeType.new(is_adjusted_to_utc: opts.fetch(:utc, true), unit: unit)) }
+      when :timestamp
+        unit = time_unit(opts.fetch(:unit, :micros))
+        converted = { millis: C::TIMESTAMP_MILLIS, micros: C::TIMESTAMP_MICROS }[unit.to_sym]
+        { type: T::INT64, converted_type: opts.fetch(:utc, true) ? converted : nil,
+          logical_type: lt(timestamp: Format::TimestampType.new(is_adjusted_to_utc: opts.fetch(:utc, true), unit: unit)) }
+      when :decimal
+        precision = Integer(opts.fetch(:precision))
+        scale = Integer(opts.fetch(:scale, 0))
+        raise ArgumentError, "Decimal precision must be positive" unless precision.positive?
+        raise ArgumentError, "Decimal scale must be between 0 and precision" unless scale.between?(0, precision)
+        physical = if opts[:physical]
+          { int32: { type: T::INT32 }, int64: { type: T::INT64 }, binary: { type: T::BYTE_ARRAY },
+            fixed: { type: T::FIXED_LEN_BYTE_ARRAY, type_length: decimal_bytes(precision) } }.fetch(opts[:physical])
+        elsif precision <= 9 then { type: T::INT32 }
+        elsif precision <= 18 then { type: T::INT64 }
+        else { type: T::FIXED_LEN_BYTE_ARRAY, type_length: decimal_bytes(precision) }
+        end
+        physical.merge(converted_type: C::DECIMAL, scale: scale, precision: precision,
+          logical_type: lt(decimal: Format::DecimalType.new(scale: scale, precision: precision)))
+      else
+        raise ArgumentError, "Unknown column type #{type.inspect}"
+      end
+    end
+
+    # Minimal number of bytes to hold a signed integer of +precision+ decimal digits
+    def decimal_bytes(precision)
+      max = 10**precision
+      n = 1
+      n += 1 while (1 << (8 * n - 1)) <= max
+      n
+    end
+
+    # Normalized logical kind of a node: [symbol, details]
+    def logical_of(node)
+      if (kind = node.logical_type&.kind)
+        name, payload = kind
+        case name
+        when :integer then return [:integer, payload.bit_width, payload.is_signed]
+        when :decimal then return [:decimal, payload.scale, payload.precision]
+        when :timestamp then return [:timestamp, payload.unit.to_sym, payload.is_adjusted_to_utc]
+        when :time then return [:time, payload.unit.to_sym, payload.is_adjusted_to_utc]
+        else return [name]
+        end
+      end
+      case node.converted_type
+      when C::UTF8 then [:string]
+      when C::ENUM then [:enum]
+      when C::JSON then [:json]
+      when C::BSON then [:bson]
+      when C::DECIMAL then [:decimal, node.scale || 0, node.precision]
+      when C::DATE then [:date]
+      when C::TIME_MILLIS then [:time, :millis, true]
+      when C::TIME_MICROS then [:time, :micros, true]
+      when C::TIMESTAMP_MILLIS then [:timestamp, :millis, true]
+      when C::TIMESTAMP_MICROS then [:timestamp, :micros, true]
+      when C::UINT_8 then [:integer, 8, false]
+      when C::UINT_16 then [:integer, 16, false]
+      when C::UINT_32 then [:integer, 32, false]
+      when C::UINT_64 then [:integer, 64, false]
+      when C::INT_8 then [:integer, 8, true]
+      when C::INT_16 then [:integer, 16, true]
+      when C::INT_32 then [:integer, 32, true]
+      when C::INT_64 then [:integer, 64, true]
+      else [nil]
+      end
+    end
+
+    # Returns a lambda converting a physical value into a Ruby value, or nil when no conversion is needed.
+    def reader_for(node)
+      kind, a, b = logical_of(node)
+      type = node.type
+      case kind
+      when :string, :enum, :json
+        return ->(v) { v.force_encoding(Encoding::UTF_8) } if type == T::BYTE_ARRAY || type == T::FIXED_LEN_BYTE_ARRAY
+      when :integer
+        if !b && (type == T::INT32 || type == T::INT64)
+          mask = type == T::INT32 ? 0xFFFF_FFFF : 0xFFFF_FFFF_FFFF_FFFF
+          return ->(v) { v & mask }
+        end
+      when :date
+        return ->(v) { Date.jd(EPOCH_JD + v, Date::GREGORIAN) } if type == T::INT32
+      when :timestamp
+        return timestamp_reader(a) if type == T::INT64
+      when :decimal
+        return decimal_reader(type, a)
+      when :uuid
+        return ->(v) { v.unpack1("H*").insert(20, "-").insert(16, "-").insert(12, "-").insert(8, "-") } if type == T::FIXED_LEN_BYTE_ARRAY
+      when :float16
+        return ->(v) { half_to_float(v.unpack1("v")) } if type == T::FIXED_LEN_BYTE_ARRAY
+      end
+      return int96_reader if type == T::INT96
+      nil
+    end
+
+    def timestamp_reader(unit)
+      div = UNIT_DIVISORS.fetch(unit)
+      name = UNIT_NAMES.fetch(unit)
+      ->(v) { Time.at(v / div, v % div, name).utc }
+    end
+
+    def int96_reader
+      lambda do |(nanos, day)|
+        secs = (day - JULIAN_EPOCH_DAY) * 86_400
+        Time.at(secs + nanos / 1_000_000_000, nanos % 1_000_000_000, :nanosecond).utc
+      end
+    end
+
+    def decimal_reader(type, scale)
+      to_decimal = scale.zero? ? ->(i) { BigDecimal(i) } : ->(i) { BigDecimal("#{i}e-#{scale}") }
+      case type
+      when T::INT32, T::INT64 then to_decimal
+      when T::BYTE_ARRAY, T::FIXED_LEN_BYTE_ARRAY
+        ->(v) { to_decimal.call(be_to_int(v)) }
+      end
+    end
+
+    # Big-endian two's complement bytes to Integer
+    def be_to_int(bytes)
+      return 0 if bytes.empty?
+      i = bytes.unpack1("H*").to_i(16)
+      bits = bytes.bytesize * 8
+      i >= (1 << (bits - 1)) ? i - (1 << bits) : i
+    end
+
+    def int_to_be(i, nbytes)
+      bits = nbytes * 8
+      raise EncodeError, "Decimal value #{i} does not fit in #{nbytes} bytes" unless i.bit_length < bits
+      i += 1 << bits if i.negative?
+      [i.to_s(16).rjust(nbytes * 2, "0")].pack("H*")
+    end
+
+    def half_to_float(h)
+      sign = (h >> 15).zero? ? 1.0 : -1.0
+      exp = (h >> 10) & 0x1F
+      frac = h & 0x3FF
+      if exp.zero?
+        sign * frac * 2.0**-24
+      elsif exp == 31
+        frac.zero? ? sign * Float::INFINITY : Float::NAN
+      else
+        sign * (1 + frac / 1024.0) * 2.0**(exp - 15)
+      end
+    end
+
+    # Rounds a Float to the nearest half-precision value (ties to even). Every step is exact
+    # in double arithmetic, so there is no double rounding.
+    def float_to_half(f)
+      return 0x7E00 if f.nan?
+      sign = f.negative? || (f.zero? && (1.0 / f).negative?) ? 0x8000 : 0
+      a = f.abs
+      return sign | 0x7C00 if a >= 65_520.0
+      return sign | (a * 2.0**24).round(half: :even) if a < 2.0**-14
+      e = Math.frexp(a)[1] - 1
+      m = ((a / 2.0**e - 1) * 1024).round(half: :even)
+      if m == 1024
+        m = 0
+        e += 1
+      end
+      sign | ((e + 15) << 10) | m
+    end
+
+    # Returns a lambda converting a Ruby value into the physical value for +node+.
+    # Besides the canonical Ruby type of each column (see the table at the top), columns accept
+    # the values Rails and plain Ruby code commonly hand over:
+    #   date:      Date, Time/DateTime (its calendar date), ISO-8601 String, Integer days since epoch
+    #   timestamp: Time, DateTime, ActiveSupport::TimeWithZone, Date (midnight UTC), ISO-8601 String,
+    #              Integer in the column's unit
+    #   time:      Time (its time of day), "HH:MM:SS[.fraction]" String, Integer in the column's unit
+    #   json:      String (used as-is) or any other object (serialized with JSON.generate)
+    #   string:    String, Symbol or anything responding to to_s
+    #   boolean:   true/false, 1/0, "true"/"false", "t"/"f", "1"/"0", "yes"/"no"
+    #   integers:  Integer, or a Float/BigDecimal/Rational/String holding a whole number
+    #   decimal:   BigDecimal, Integer, Rational, Float or numeric String
+    #   uuid:      String with or without dashes, or 16 raw bytes
+    def writer_for(node)
+      kind, a, b = logical_of(node)
+      type = node.type
+      case kind
+      when :date
+        return method(:date_to_days)
+      when :timestamp
+        return timestamp_writer(a, b)
+      when :time
+        return time_writer(a) if type == T::INT32 || type == T::INT64
+      when :json
+        return ->(v) { v.is_a?(String) ? v : JSON.generate(v) }
+      when :string, :enum
+        values = node.enum_values
+        return enum_writer(values) if values
+        return ->(v) { v.is_a?(String) ? v : v.to_s }
+      when :decimal
+        return decimal_writer(node, a)
+      when :uuid
+        return fixed_checker(16) do |v|
+          v.bytesize == 16 && v.encoding == Encoding::BINARY ? v : [v.to_s.delete("-")].pack("H*")
+        end
+      when :float16
+        return ->(v) { [float_to_half(Float(v))].pack("v") }
+      when :integer
+        if !b && (type == T::INT32 || type == T::INT64)
+          bits = type == T::INT32 ? 32 : 64
+          check = int_checker(0, (1 << a) - 1)
+          return ->(v) { Delta.wrap(check.call(v), bits) }
+        elsif type == T::INT32 || type == T::INT64
+          return int_checker(-(1 << (a - 1)), (1 << (a - 1)) - 1)
+        end
+      end
+
+      case type
+      when T::BOOLEAN then method(:to_boolean)
+      when T::INT32 then int_checker(-(1 << 31), (1 << 31) - 1)
+      when T::INT64 then int_checker(-(1 << 63), (1 << 63) - 1)
+      when T::FLOAT, T::DOUBLE then ->(v) { v.is_a?(Float) ? v : Float(v) }
+      when T::BYTE_ARRAY then ->(v) { v.is_a?(String) ? v : v.to_s }
+      when T::FIXED_LEN_BYTE_ARRAY then fixed_checker(node.type_length) { |v| v.is_a?(String) ? v : v.to_s }
+      when T::INT96
+        lambda do |v|
+          return v if v.is_a?(Array)
+          v = to_time(v)
+          nanos = v.to_i * 1_000_000_000 + v.nsec
+          day, nanos_of_day = nanos.divmod(NANOS_PER_DAY)
+          [nanos_of_day, day + JULIAN_EPOCH_DAY]
+        end
+      end
+    end
+
+    BOOLEANS = {
+      true => true, false => false, 1 => true, 0 => false,
+      "true" => true, "false" => false, "t" => true, "f" => false, "1" => true, "0" => false,
+      "yes" => true, "no" => false, "TRUE" => true, "FALSE" => false, "T" => true, "F" => false
+    }.freeze
+
+    def to_boolean(v)
+      BOOLEANS.fetch(v) do
+        s = v.is_a?(String) || v.is_a?(Symbol) ? v.to_s.downcase : nil
+        BOOLEANS.fetch(s) { raise ArgumentError, "expected a boolean, got #{v.inspect}" }
+      end
+    end
+
+    def date_to_days(v)
+      return v if v.is_a?(Integer)
+      d = case v
+      when Date then v
+      when String then Date.iso8601(v)
+      else
+        raise ArgumentError, "expected a Date, got #{v.class}" unless v.respond_to?(:to_date)
+        v.to_date
+      end
+      # Use the civil date, so dates before 1582 in Ruby's default calendar are proleptic Gregorian
+      Date.civil(d.year, d.mon, d.mday, Date::GREGORIAN).jd - EPOCH_JD
+    end
+
+    # Converts the values a timestamp column accepts into a Time (or Time-like) object
+    def to_time(v)
+      case v
+      when Time then v
+      when DateTime then v.to_time
+      when Date then Time.utc(v.year, v.month, v.day)
+      when String
+        begin
+          Time.iso8601(v)
+        rescue ArgumentError
+          Time.parse(v)
+        end
+      else
+        # ActiveSupport::TimeWithZone and friends
+        raise ArgumentError, "expected a Time, got #{v.class}" unless v.respond_to?(:to_i) && v.respond_to?(:nsec)
+        v
+      end
+    end
+
+    def timestamp_writer(unit, utc)
+      mult = UNIT_DIVISORS.fetch(unit)
+      lambda do |v|
+        return v if v.is_a?(Integer)
+        t = to_time(v)
+        secs = t.to_i
+        # Local (not UTC-adjusted) timestamps store the wall clock time
+        secs += t.utc_offset unless utc
+        secs * mult + t.nsec * mult / 1_000_000_000
+      end
+    end
+
+    def time_writer(unit)
+      mult = UNIT_DIVISORS.fetch(unit)
+      lambda do |v|
+        return v if v.is_a?(Integer)
+        if v.is_a?(String)
+          m = /\A(\d{1,2}):(\d{2})(?::(\d{2})(?:\.(\d{1,9}))?)?\z/.match(v) or
+            raise ArgumentError, "expected HH:MM[:SS[.fraction]], got #{v.inspect}"
+          nanos = (m[4] || "").ljust(9, "0").to_i
+          secs = m[1].to_i * 3600 + m[2].to_i * 60 + m[3].to_i
+        else
+          raise ArgumentError, "expected a Time, got #{v.class}" unless v.respond_to?(:hour) && v.respond_to?(:nsec)
+          nanos = v.nsec
+          secs = v.hour * 3600 + v.min * 60 + v.sec
+        end
+        raise RangeError, "time of day out of range: #{v.inspect}" unless secs < 86_400
+        secs * mult + nanos * mult / 1_000_000_000
+      end
+    end
+
+    # String column restricted to a set of values. +values+ is an Array of labels, or a Hash of
+    # label => stored value like Rails' `Model.statuses`, in which case either is accepted.
+    def enum_writer(values)
+      labels = {}
+      if values.is_a?(Hash)
+        values.each do |label, stored|
+          labels[label.to_s] = label.to_s
+          labels[stored] = label.to_s unless stored.is_a?(String) || stored.is_a?(Symbol)
+        end
+      else
+        values.each { |label| labels[label.to_s] = label.to_s }
+      end
+      lambda do |v|
+        labels.fetch(v.is_a?(Symbol) ? v.to_s : v) do
+          raise ArgumentError, "#{v.inspect} is not one of #{labels.values.uniq.join(", ")}"
+        end
+      end
+    end
+
+    # Converts to Integer, rejecting fractional numbers and values outside min..max
+    def int_checker(min, max)
+      lambda do |v|
+        i = Integer(v)
+        raise ArgumentError, "#{v.inspect} is not an integer" unless v.is_a?(Integer) || !v.is_a?(Numeric) || v == i
+        raise RangeError, "#{i} is outside #{min}..#{max}" unless i >= min && i <= max
+        i
+      end
+    end
+
+    def fixed_checker(length, &convert)
+      lambda do |v|
+        s = convert.call(v)
+        raise ArgumentError, "expected #{length} bytes, got #{s.bytesize}" unless s.bytesize == length
+        s
+      end
+    end
+
+    def decimal_writer(node, scale)
+      mult = 10**scale
+      limit = node.precision ? 10**node.precision : nil
+      to_unscaled = lambda do |v|
+        r = case v
+        when Integer then v * mult
+        when Float, String then Rational(v.to_s) * mult
+        else v.to_r * mult
+        end
+        r = r.round if r.is_a?(Rational)
+        i = r.to_i
+        raise RangeError, "#{v} does not fit DECIMAL(#{node.precision}, #{scale})" if limit && i.abs >= limit
+        i
+      end
+      case node.type
+      when T::INT32, T::INT64 then to_unscaled
+      when T::FIXED_LEN_BYTE_ARRAY
+        len = node.type_length
+        ->(v) { int_to_be(to_unscaled.call(v), len) }
+      when T::BYTE_ARRAY
+        lambda do |v|
+          i = to_unscaled.call(v)
+          int_to_be(i, [(i.bit_length + 8) / 8, 1].max)
+        end
+      end
+    end
+  end
+end

data/lib/herringbone/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Herringbone
+  VERSION = "0.1.0"
+end