red-arrow-format 25.0.1 → 26.0.0

data/lib/arrow-format/array.rb

@@ -1,3 +1,4 @@
+# Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
 # distributed with this work for additional information
 # regarding copyright ownership.  The ASF licenses this file
@@ -15,11 +16,24 @@
 # under the License.
 
 require "bigdecimal"
+require "date"
 
+require_relative "array-builder"
 require_relative "bitmap"
+require_relative "bitmap-builder"
 
 module ArrowFormat
+  using FlatBuffers::AppendAsBytes if FlatBuffers.const_defined?(:AppendAsBytes)
+
   class Array
+    include Enumerable
+
+    class << self
+      def build(values)
+        ArrayBuilder.new(values).build
+      end
+    end
+
     attr_reader :type
     attr_reader :size
     alias_method :length, :size
@@ -48,6 +62,14 @@ module ArrowFormat
       not valid?(i)
     end
 
+    def [](i)
+      if valid?(i)
+        unpack_value(i)
+      else
+        nil
+      end
+    end
+
     def n_nulls
       if @validity_buffer.nil?
         0
@@ -60,6 +82,16 @@ module ArrowFormat
       @size.zero?
     end
 
+    def ==(other)
+      return false unless other.is_a?(self.class)
+      return false unless @size == other.size
+      return true if @validity_buffer.nil? and other.validity_buffer.nil?
+      if @offset == other.offset and @validity_buffer == other.validity_buffer
+        return true
+      end
+      validity_bitmap == other.validity_bitmap
+    end
+
     protected
     def slice!(offset, size)
       @offset = offset
@@ -67,11 +99,11 @@ module ArrowFormat
       clear_cache
     end
 
-    private
     def validity_bitmap
       @validity_bitmap ||= Bitmap.new(@validity_buffer, @offset, @size)
     end
 
+    private
     def apply_validity(array)
       return array if @validity_buffer.nil?
       validity_bitmap.each_with_index do |is_valid, i|
@@ -144,6 +176,14 @@ module ArrowFormat
       super(NullType.singleton, size, nil)
     end
 
+    def valid?(i)
+      false
+    end
+
+    def null?(i)
+      true
+    end
+
     def each_buffer
       return to_enum(__method__) unless block_given?
     end
@@ -155,10 +195,40 @@ module ArrowFormat
     def to_a
       [nil] * @size
     end
+
+    def each
+      return to_enum(__method__) unless block_given?
+
+      @size.times do
+        yield(nil)
+      end
+    end
   end
 
   class PrimitiveArray < Array
-    def initialize(type, size, validity_buffer, values_buffer)
+    include BufferAlignable
+
+    attr_reader :values_buffer
+    def initialize(*args)
+      n_args = args.size
+      if self.class.respond_to?(:type)
+        type = self.class.type
+        expected_n_args = "1 or 3"
+      else
+        type = args.shift
+        unless type.is_a?(Type)
+          type = self.class.type_class.try_convert(type) || type
+        end
+        expected_n_args = "2 or 4"
+      end
+      args = build_data(args[0], type) if args.size == 1
+      if args.size != 3
+        message =
+          "wrong number of arguments " +
+          "(given #{n_args}, expected #{expected_n_args})"
+        raise ArgumentError, message
+      end
+      size, validity_buffer, values_buffer = args
       super(type, size, validity_buffer)
       @values_buffer = values_buffer
     end
@@ -170,8 +240,30 @@ module ArrowFormat
       apply_validity(@values_buffer.values(@type.buffer_type, offset, @size))
     end
 
+    def each(&block)
+      return to_enum(__method__) {@size} unless block_given?
+
+      each_value = Enumerator.new(@size) do |yielder|
+        offset = element_size * @offset
+        @values_buffer.each(@type.buffer_type, offset, @size) do |_, value|
+          yielder << value
+        end
+      end
+      if @validity_buffer.nil?
+        each_value.each(&block)
+      else
+        validity_bitmap.zip(each_value) do |is_valid, value|
+          if is_valid
+            yield(value)
+          else
+            yield(nil)
+          end
+        end
+      end
+    end
+
     def each_buffer
-      return to_enum(__method__) unless block_given?
+      return to_enum(__method__) {2} unless block_given?
 
       yield(slice_bitmap_buffer(:validity, @validity_buffer))
       yield(slice_fixed_element_size_buffer(:values,
@@ -179,43 +271,126 @@ module ArrowFormat
                                             element_size))
     end
 
+    def ==(other)
+      return false unless super(other)
+      if @offset == other.offset and @values_buffer == other.values_buffer
+        return true
+      end
+      lazy.zip(other).all? do |value, other_value|
+        value == other_value
+      end
+    end
+
     private
     def element_size
       IO::Buffer.size_of(@type.buffer_type)
     end
+
+    def build_data(data, type)
+      n = 0
+      validity_buffer_builder = nil
+      buffer = +"".b
+      pack_template = type.pack_template
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+          buffer.append_as_bytes(pack_value(nil, pack_template, type))
+        else
+          buffer.append_as_bytes(pack_value(value, pack_template, type))
+        end
+        n += 1
+      end
+      validity_buffer = validity_buffer_builder&.finish(n)
+      pad!(buffer, buffer_padding_size(buffer))
+      buffer.freeze
+      return n, validity_buffer, IO::Buffer.for(buffer)
+    end
+
+    def pack_value(value, template, type)
+      value = 0 if value.nil?
+      [value].pack(template)
+    end
+
+    def unpack_value(i)
+      offset = element_size * (@offset + i)
+      @values_buffer.get_value(@type.buffer_type, offset)
+    end
   end
 
   class BooleanArray < PrimitiveArray
-    def initialize(size, validity_buffer, values_buffer)
-      super(BooleanType.singleton, size, validity_buffer, values_buffer)
+    class << self
+      def type
+        BooleanType.singleton
+      end
     end
 
     def to_a
       return [] if empty?
 
-      @values_bitmap ||= Bitmap.new(@values_buffer, @offset, @size)
-      values = @values_bitmap.to_a
+      values = values_bitmap.to_a
       apply_validity(values)
     end
 
+    def each(&block)
+      return to_enum(__method__) {@size} unless block_given?
+
+      if @validity_buffer.nil?
+        values_bitmap.each(&block)
+      else
+        validity_bitmap.zip(values_bitmap) do |is_valid, value|
+          if is_valid
+            yield(value)
+          else
+            yield(nil)
+          end
+        end
+      end
+    end
+
     def each_buffer
-      return to_enum(__method__) unless block_given?
+      return to_enum(__method__) {2} unless block_given?
 
       yield(slice_bitmap_buffer(:validity, @validity_buffer))
       yield(slice_bitmap_buffer(:values, @values_buffer))
     end
 
     private
+    def values_bitmap
+      @values_bitmap ||= Bitmap.new(@values_buffer, @offset, @size)
+    end
+
     def clear_cache
       super
       @values_bitmap = nil
     end
+
+    def build_data(data, type)
+      n = 0
+      validity_buffer_builder = nil
+      values_buffer_builder = DenseBitmapBuilder.new
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+          values_buffer_builder.append(false)
+        elsif value
+          values_buffer_builder.append(true)
+        else
+          values_buffer_builder.append(false)
+        end
+        n += 1
+      end
+      validity_buffer = validity_buffer_builder&.finish(n)
+      return n, validity_buffer, values_buffer_builder.finish
+    end
+
+    def unpack_value(i)
+      values_bitmap[i]
+    end
   end
 
   class IntArray < PrimitiveArray
-    def initialize(size, validity_buffer, values_buffer)
-      super(self.class.type, size, validity_buffer, values_buffer)
-    end
   end
 
   class Int8Array < IntArray
@@ -283,9 +458,6 @@ module ArrowFormat
   end
 
   class FloatingPointArray < PrimitiveArray
-    def initialize(size, validity_buffer, values_buffer)
-      super(self.class.type, size, validity_buffer, values_buffer)
-    end
   end
 
   class Float32Array < FloatingPointArray
@@ -308,9 +480,6 @@ module ArrowFormat
   end
 
   class DateArray < TemporalArray
-    def initialize(size, validity_buffer, values_buffer)
-      super(self.class.type, size, validity_buffer, values_buffer)
-    end
   end
 
   class Date32Array < DateArray
@@ -319,6 +488,18 @@ module ArrowFormat
         Date32Type.singleton
       end
     end
+
+    private
+    UNIX_EPOCH = Date.new(1970, 1, 1).jd
+    def pack_value(value, template, type)
+      if value.nil?
+        [0].pack(template)
+      elsif value.is_a?(Date)
+        [value.jd - UNIX_EPOCH].pack(template)
+      else
+        [value].pack(template)
+      end
+    end
   end
 
   class Date64Array < DateArray
@@ -333,21 +514,68 @@ module ArrowFormat
   end
 
   class Time32Array < TimeArray
+    class << self
+      def type_class
+        Time32Type
+      end
+    end
   end
 
   class Time64Array < TimeArray
+    class << self
+      def type_class
+        Time64Type
+      end
+    end
   end
 
   class TimestampArray < TemporalArray
+    class << self
+      def type_class
+        TimestampType
+      end
+    end
+
+    private
+    def pack_value(value, template, type)
+      if value.nil?
+        [0, 0].pack(template)
+      elsif value.is_a?(Time)
+        case type.unit
+        when :second
+          value = value.to_i
+        when :millisecond
+          value = (value.to_i * 1_000) + (value.nsec / 1_000_000)
+        when :microsecond
+          value = (value.to_i * 1_000_000) + (value.nsec / 1_000)
+        when :nanosecond
+          value = (value.to_i * 1_000_000_000) + value.nsec
+        end
+        [value].pack(template)
+      else
+        [value].pack(template)
+      end
+    end
   end
 
   class IntervalArray < TemporalArray
   end
 
   class YearMonthIntervalArray < IntervalArray
+    class << self
+      def type
+        YearMonthIntervalType.singleton
+      end
+    end
   end
 
   class DayTimeIntervalArray < IntervalArray
+    class << self
+      def type
+        DayTimeIntervalType.singleton
+      end
+    end
+
     def to_a
       return [] if empty?
 
@@ -355,19 +583,64 @@ module ArrowFormat
       values = @values_buffer.
                  each(@type.buffer_type, offset, @size * 2).
                  each_slice(2).
-                 collect do |(_, day), (_, time)|
-        [day, time]
+                 collect do |(_, day), (_, millisecond)|
+        [day, millisecond]
       end
       apply_validity(values)
     end
 
+    def each(&block)
+      return to_enum(__method__) {@size} unless block_given?
+
+      each_value = Enumerator.new(@size) do |yielder|
+        offset = element_size * @offset
+        @values_buffer.
+          each(@type.buffer_type, offset, @size * 2).
+          each_slice(2) do |(_, day), (_, millisecond)|
+          yielder << [day, millisecond]
+        end
+      end
+      if @validity_buffer.nil?
+        each_value.each(&block)
+      else
+        validity_bitmap.zip(each_value) do |is_valid, value|
+          if is_valid
+            yield(value)
+          else
+            yield(nil)
+          end
+        end
+      end
+    end
+
     private
     def element_size
       super * 2
     end
+
+    def pack_value(value, template, type)
+      if value.nil?
+        [0, 0].pack(template)
+      elsif value.is_a?(Hash)
+        [value[:day], value[:millisecond]].pack(template)
+      else
+        value.pack(template)
+      end
+    end
+
+    def unpack_value(i)
+      offset = element_size * (@offset + i)
+      @values_buffer.get_values([@type.buffer_type] * 2, offset)
+    end
   end
 
   class MonthDayNanoIntervalArray < IntervalArray
+    class << self
+      def type
+        MonthDayNanoIntervalType.singleton
+      end
+    end
+
     def to_a
       return [] if empty?
 
@@ -381,17 +654,75 @@ module ArrowFormat
       apply_validity(values)
     end
 
+    def each(&block)
+      return to_enum(__method__) {@size} unless block_given?
+
+      each_value = Enumerator.new(@size) do |yielder|
+        buffer_types = @type.buffer_types
+        value_size = IO::Buffer.size_of(buffer_types)
+        base_offset = value_size * @offset
+        @size.times do |i|
+          offset = base_offset + value_size * i
+          yielder << @values_buffer.get_values(buffer_types, offset)
+        end
+      end
+      if @validity_buffer.nil?
+        each_value.each(&block)
+      else
+        validity_bitmap.zip(each_value) do |is_valid, value|
+          if is_valid
+            yield(value)
+          else
+            yield(nil)
+          end
+        end
+      end
+    end
+
     private
     def element_size
       IO::Buffer.size_of(@type.buffer_types)
     end
+
+    def pack_value(value, template, type)
+      if value.nil?
+        [0, 0, 0].pack(template)
+      elsif value.is_a?(Hash)
+        [value[:month], value[:day], value[:nanosecond]].pack(template)
+      else
+        value.pack(template)
+      end
+    end
+
+    def unpack_value(i)
+      buffer_types = @type.buffer_types
+      value_size = IO::Buffer.size_of(buffer_types)
+      offset = value_size * (@offset + i)
+      @values_buffer.get_values(buffer_types, offset)
+    end
   end
 
   class DurationArray < TemporalArray
+    class << self
+      def type_class
+        DurationType
+      end
+    end
   end
 
   class VariableSizeBinaryArray < Array
-    def initialize(size, validity_buffer, offsets_buffer, values_buffer)
+    include BufferAlignable
+
+    def initialize(*args)
+      if args.size == 1
+        args = build_data(args.first)
+      elsif args.size != 4
+        raise ArgumentError,
+              "wrong number of arguments (given #{args.size}, expected 1 or 4)"
+      end
+
+      size, validity_buffer, offsets_buffer, values_buffer = args
+
       super(self.class.type, size, validity_buffer)
       @offsets_buffer = offsets_buffer
       @values_buffer = values_buffer
@@ -431,9 +762,60 @@ module ArrowFormat
     end
 
     private
+    def build_data(data)
+      type = self.class.type
+
+      n = 0
+      validity_buffer_builder = nil
+
+      values = +"".b
+      offsets = [0]
+
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+
+          offsets << values.bytesize
+        else
+          values.append_as_bytes(value)
+          offsets << values.bytesize
+        end
+
+        n += 1
+      end
+
+      validity_buffer = validity_buffer_builder&.finish(n)
+
+      offsets_data = offsets.pack("#{type.offset_pack_template}*")
+      pad!(offsets_data, buffer_padding_size(offsets_data))
+      offsets_data.freeze
+      offsets_buffer = IO::Buffer.for(offsets_data)
+
+      pad!(values, buffer_padding_size(values))
+      values.freeze
+      values_buffer = IO::Buffer.for(values)
+
+      [
+        n,
+        validity_buffer,
+        offsets_buffer,
+        values_buffer,
+      ]
+    end
+
     def offset_size
       IO::Buffer.size_of(@type.offset_buffer_type)
     end
+
+    def unpack_value(i)
+      offset_types = [@type.offset_buffer_type] * 2
+      offsets_offset = offset_size * (@offset + i)
+      offset, next_offset = @offsets_buffer.get_values(offset_types,
+                                                       offsets_offset)
+      length = next_offset - offset
+      @values_buffer.get_string(offset, length, @type.encoding)
+    end
   end
 
   class BinaryArray < VariableSizeBinaryArray
@@ -472,7 +854,21 @@ module ArrowFormat
   end
 
   class FixedSizeBinaryArray < Array
-    def initialize(type, size, validity_buffer, values_buffer)
+    include BufferAlignable
+
+    def initialize(type, *args)
+      unless type.is_a?(Type)
+        type = FixedSizeBinaryType.try_convert(type) || type
+      end
+      if args.size == 1
+        args = build_data(args.first, type)
+      elsif args.size != 3
+        raise ArgumentError,
+              "wrong number of arguments (given #{args.size + 1}, expected 2 or 4)"
+      end
+
+      size, validity_buffer, values_buffer = args
+
       super(type, size, validity_buffer)
       @values_buffer = values_buffer
     end
@@ -495,6 +891,44 @@ module ArrowFormat
       end
       apply_validity(values)
     end
+
+    private
+    def build_data(data, type)
+      n = 0
+      validity_buffer_builder = nil
+
+      values = +"".b
+      byte_width = type.byte_width
+      null_value = "\x00" * byte_width
+
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+          values.append_as_bytes(null_value)
+        else
+          unless value.bytesize == byte_width
+            message = "value size must be #{byte_width}: #{value.inspect}"
+            raise ArgumentError, message
+          end
+          values.append_as_bytes(value)
+        end
+
+        n += 1
+      end
+
+      validity_buffer = validity_buffer_builder&.finish(n)
+
+      pad!(values, buffer_padding_size(values))
+      values.freeze
+      values_buffer = IO::Buffer.for(values)
+
+      [
+        n,
+        validity_buffer,
+        values_buffer,
+      ]
+    end
   end
 
   class DecimalArray < FixedSizeBinaryArray
@@ -605,6 +1039,53 @@ module ArrowFormat
   end
 
   class ListArray < VariableSizeListArray
+    include BufferAlignable
+
+    def initialize(type, *args)
+      if args.size == 1
+        args = build_data(type, args.first)
+      elsif args.size != 4
+        raise ArgumentError,
+              "wrong number of arguments (given #{args.size + 1}, expected 2 or 5)"
+      end
+
+      super(type, *args)
+    end
+
+    private
+    def build_data(type, data)
+      n = 0
+      validity_buffer_builder = nil
+
+      child_values = []
+      offsets = [0]
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+        else
+          child_values.concat(value)
+        end
+        offsets << child_values.size
+        n += 1
+      end
+
+      validity_buffer = validity_buffer_builder&.finish(n)
+
+      offsets_data = offsets.pack("#{type.offset_pack_template}*")
+      pad!(offsets_data, buffer_padding_size(offsets_data))
+      offsets_data.freeze
+      offsets_buffer = IO::Buffer.for(offsets_data)
+
+      child = type.child.type.build_array(child_values)
+
+      [
+        n,
+        validity_buffer,
+        offsets_buffer,
+        child,
+      ]
+    end
   end
 
   class LargeListArray < VariableSizeListArray
@@ -612,7 +1093,15 @@ module ArrowFormat
 
   class FixedSizeListArray < Array
     attr_reader :child
-    def initialize(type, size, validity_buffer, child)
+    def initialize(type, *args)
+      if args.size == 1
+        args = build_data(type, args.first)
+      elsif args.size != 3
+        raise ArgumentError,
+              "wrong number of arguments (given #{args.size + 1}, expected 2 or 4)"
+      end
+
+      size, validity_buffer, child = args
       super(type, size, validity_buffer)
       @child = child
     end
@@ -631,6 +1120,36 @@ module ArrowFormat
     end
 
     private
+    def build_data(type, data)
+      n = 0
+      validity_buffer_builder = nil
+
+      child_values = []
+      data.each_with_index do |value, i|
+        if value.nil?
+          validity_buffer_builder ||= SparseBitmapBuilder.new
+          validity_buffer_builder.unset(i)
+          child_values.concat([nil] * type.size)
+        else
+          unless value.size == type.size
+            message = "list size must be #{type.size}: #{value.inspect}"
+            raise ArgumentError, message
+          end
+          child_values.concat(value)
+        end
+        n += 1
+      end
+
+      validity_buffer = validity_buffer_builder&.finish(n)
+      child = type.child.type.build_array(child_values)
+
+      [
+        n,
+        validity_buffer,
+        child,
+      ]
+    end
+
     def slice!(offset, size)
       super
       @child = @child.slice(@type.size * @offset,