red-arrow-format 23.0.0 → 24.0.0

data/lib/arrow-format/type.rb

@@ -16,6 +16,9 @@
 
 module ArrowFormat
   class Type
+    def to_s
+      name
+    end
   end
 
   class NullType < Type
@@ -30,11 +33,18 @@ module ArrowFormat
     end
 
     def build_array(size)
-      NullArray.new(self, size)
+      NullArray.new(size)
+    end
+
+    def to_flatbuffers
+      FB::Null::Data.new
     end
   end
 
-  class BooleanType < Type
+  class PrimitiveType < Type
+  end
+
+  class BooleanType < PrimitiveType
     class << self
       def singleton
         @singleton ||= new
@@ -46,21 +56,35 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      BooleanArray.new(self, size, validity_buffer, values_buffer)
+      BooleanArray.new(size, validity_buffer, values_buffer)
+    end
+
+    def to_flatbuffers
+      FB::Bool::Data.new
     end
   end
 
-  class NumberType < Type
+  class NumberType < PrimitiveType
   end
 
   class IntType < NumberType
     attr_reader :bit_width
-    attr_reader :signed
     def initialize(bit_width, signed)
       super()
       @bit_width = bit_width
       @signed = signed
     end
+
+    def signed?
+      @signed
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Int::Data.new
+      fb_type.bit_width = @bit_width
+      fb_type.signed = @signed
+      fb_type
+    end
   end
 
   class Int8Type < IntType
@@ -83,7 +107,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      Int8Array.new(self, size, validity_buffer, values_buffer)
+      Int8Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -107,7 +131,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      UInt8Array.new(self, size, validity_buffer, values_buffer)
+      UInt8Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -131,7 +155,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      Int16Array.new(self, size, validity_buffer, values_buffer)
+      Int16Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -155,7 +179,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      UInt16Array.new(self, size, validity_buffer, values_buffer)
+      UInt16Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -179,7 +203,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      Int32Array.new(self, size, validity_buffer, values_buffer)
+      Int32Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -203,7 +227,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      UInt32Array.new(self, size, validity_buffer, values_buffer)
+      UInt32Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -227,7 +251,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      Int64Array.new(self, size, validity_buffer, values_buffer)
+      Int64Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -251,7 +275,7 @@ module ArrowFormat
     end
 
     def build_array(size, validity_buffer, values_buffer)
-      UInt64Array.new(self, size, validity_buffer, values_buffer)
+      UInt64Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -261,6 +285,12 @@ module ArrowFormat
       super()
       @precision = precision
     end
+
+    def to_flatbuffers
+      fb_type = FB::FloatingPoint::Data.new
+      fb_type.precision = FB::Precision.try_convert(@precision.to_s.upcase)
+      fb_type
+    end
   end
 
   class Float32Type < FloatingPointType
@@ -278,8 +308,12 @@ module ArrowFormat
       "Float32"
     end
 
+    def buffer_type
+      :f32
+    end
+
     def build_array(size, validity_buffer, values_buffer)
-      Float32Array.new(self, size, validity_buffer, values_buffer)
+      Float32Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -298,15 +332,30 @@ module ArrowFormat
       "Float64"
     end
 
+    def buffer_type
+      :f64
+    end
+
     def build_array(size, validity_buffer, values_buffer)
-      Float64Array.new(self, size, validity_buffer, values_buffer)
+      Float64Array.new(size, validity_buffer, values_buffer)
     end
   end
 
-  class TemporalType < Type
+  class TemporalType < PrimitiveType
   end
 
   class DateType < TemporalType
+    attr_reader :unit
+    def initialize(unit)
+      super()
+      @unit = unit
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Date::Data.new
+      fb_type.unit = FB::DateUnit.try_convert(@unit.to_s.upcase)
+      fb_type
+    end
   end
 
   class Date32Type < DateType
@@ -316,12 +365,20 @@ module ArrowFormat
       end
     end
 
+    def initialize
+      super(:day)
+    end
+
     def name
       "Date32"
     end
 
+    def buffer_type
+      :s32
+    end
+
     def build_array(size, validity_buffer, values_buffer)
-      Date32Array.new(self, size, validity_buffer, values_buffer)
+      Date32Array.new(size, validity_buffer, values_buffer)
     end
   end
 
@@ -332,38 +389,75 @@ module ArrowFormat
       end
     end
 
+    def initialize
+      super(:millisecond)
+    end
+
     def name
       "Date64"
     end
 
+    def buffer_type
+      :s64
+    end
+
     def build_array(size, validity_buffer, values_buffer)
-      Date64Array.new(self, size, validity_buffer, values_buffer)
+      Date64Array.new(size, validity_buffer, values_buffer)
     end
   end
 
   class TimeType < TemporalType
+    attr_reader :bit_width
     attr_reader :unit
-    def initialize(unit)
+    def initialize(bit_width, unit)
       super()
+      @bit_width = bit_width
       @unit = unit
     end
+
+    def to_s
+      "#{super}(#{unit})"
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Time::Data.new
+      fb_type.bit_width = @bit_width
+      fb_type.unit = FB::TimeUnit.try_convert(@unit.to_s.upcase)
+      fb_type
+    end
   end
 
   class Time32Type < TimeType
+    def initialize(unit)
+      super(32, unit)
+    end
+
     def name
       "Time32"
     end
 
+    def buffer_type
+      :s32
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       Time32Array.new(self, size, validity_buffer, values_buffer)
     end
   end
 
   class Time64Type < TimeType
+    def initialize(unit)
+      super(64, unit)
+    end
+
     def name
       "Time64"
     end
 
+    def buffer_type
+      :s64
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       Time64Array.new(self, size, validity_buffer, values_buffer)
     end
@@ -371,50 +465,108 @@ module ArrowFormat
 
   class TimestampType < TemporalType
     attr_reader :unit
-    attr_reader :timezone
-    def initialize(unit, timezone)
+    attr_reader :time_zone
+    def initialize(unit, time_zone)
       super()
       @unit = unit
-      @timezone = timezone
+      @time_zone = time_zone
     end
 
     def name
       "Timestamp"
     end
 
+    def buffer_type
+      :s64
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       TimestampArray.new(self, size, validity_buffer, values_buffer)
     end
+
+    def to_s
+      options = [@unit]
+      options << @time_zone if @time_zone
+      "#{super}(#{options.join(", ")})"
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Timestamp::Data.new
+      fb_type.unit = FB::TimeUnit.try_convert(@unit.to_s.upcase)
+      fb_type.timezone = @time_zone
+      fb_type
+    end
   end
 
   class IntervalType < TemporalType
+    class << self
+      def singleton
+        @singleton ||= new
+      end
+    end
+
+    attr_reader :unit
+    def initialize(unit)
+      super()
+      @unit = unit
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Interval::Data.new
+      fb_type.unit = FB::IntervalUnit.try_convert(@unit.to_s.upcase)
+      fb_type
+    end
   end
 
   class YearMonthIntervalType < IntervalType
+    def initialize
+      super(:year_month)
+    end
+
     def name
       "YearMonthInterval"
     end
 
+    def buffer_type
+      :s32
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       YearMonthIntervalArray.new(self, size, validity_buffer, values_buffer)
     end
   end
 
   class DayTimeIntervalType < IntervalType
+    def initialize
+      super(:day_time)
+    end
+
     def name
       "DayTimeInterval"
     end
 
+    def buffer_type
+      :s32
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       DayTimeIntervalArray.new(self, size, validity_buffer, values_buffer)
     end
   end
 
   class MonthDayNanoIntervalType < IntervalType
+    def initialize
+      super(:month_day_nano)
+    end
+
     def name
       "MonthDayNanoInterval"
     end
 
+    def buffer_types
+      @buffer_types ||= [:s32, :s32, :s64]
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       MonthDayNanoIntervalArray.new(self,
                                     size,
@@ -434,9 +586,23 @@ module ArrowFormat
       "Duration"
     end
 
+    def buffer_type
+      :s64
+    end
+
     def build_array(size, validity_buffer, values_buffer)
       DurationArray.new(self, size, validity_buffer, values_buffer)
     end
+
+    def to_s
+      "#{super}(#{@unit})"
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Duration::Data.new
+      fb_type.unit = FB::TimeUnit.try_convert(@unit.to_s.upcase)
+      fb_type
+    end
   end
 
   class VariableSizeBinaryType < Type
@@ -453,8 +619,23 @@ module ArrowFormat
       "Binary"
     end
 
+    def offset_buffer_type
+      :s32 # TODO: big endian support
+    end
+
+    def encoding
+      Encoding::ASCII_8BIT
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, values_buffer)
-      BinaryArray.new(self, size, validity_buffer, offsets_buffer, values_buffer)
+      BinaryArray.new(size,
+                      validity_buffer,
+                      offsets_buffer,
+                      values_buffer)
+    end
+
+    def to_flatbuffers
+      FB::Binary::Data.new
     end
   end
 
@@ -469,13 +650,24 @@ module ArrowFormat
       "LargeBinary"
     end
 
+    def offset_buffer_type
+      :s64 # TODO: big endian support
+    end
+
+    def encoding
+      Encoding::ASCII_8BIT
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, values_buffer)
-      LargeBinaryArray.new(self,
-                           size,
+      LargeBinaryArray.new(size,
                            validity_buffer,
                            offsets_buffer,
                            values_buffer)
     end
+
+    def to_flatbuffers
+      FB::LargeBinary::Data.new
+    end
   end
 
   class UTF8Type < VariableSizeBinaryType
@@ -489,8 +681,20 @@ module ArrowFormat
       "UTF8"
     end
 
+    def offset_buffer_type
+      :s32 # TODO: big endian support
+    end
+
+    def encoding
+      Encoding::UTF_8
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, values_buffer)
-      UTF8Array.new(self, size, validity_buffer, offsets_buffer, values_buffer)
+      UTF8Array.new(size, validity_buffer, offsets_buffer, values_buffer)
+    end
+
+    def to_flatbuffers
+      FB::Utf8::Data.new
     end
   end
 
@@ -505,13 +709,24 @@ module ArrowFormat
       "LargeUTF8"
     end
 
+    def offset_buffer_type
+      :s64 # TODO: big endian support
+    end
+
+    def encoding
+      Encoding::UTF_8
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, values_buffer)
-      LargeUTF8Array.new(self,
-                         size,
+      LargeUTF8Array.new(size,
                          validity_buffer,
                          offsets_buffer,
                          values_buffer)
     end
+
+    def to_flatbuffers
+      FB::LargeUtf8::Data.new
+    end
   end
 
   class FixedSizeBinaryType < Type
@@ -528,6 +743,16 @@ module ArrowFormat
     def build_array(size, validity_buffer, values_buffer)
       FixedSizeBinaryArray.new(self, size, validity_buffer, values_buffer)
     end
+
+    def to_s
+      "#{super}(#{@byte_width})"
+    end
+
+    def to_flatbuffers
+      fb_type = FB::FixedSizeBinary::Data.new
+      fb_type.byte_width = @byte_width
+      fb_type
+    end
   end
 
   class DecimalType < FixedSizeBinaryType
@@ -538,6 +763,18 @@ module ArrowFormat
       @precision = precision
       @scale = scale
     end
+
+    def to_s
+      "#{name}(#{@precision}, #{@scale})"
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Decimal::Data.new
+      fb_type.bit_width = @byte_width * 8
+      fb_type.precision = @precision
+      fb_type.scale = @scale
+      fb_type
+    end
   end
 
   class Decimal128Type < DecimalType
@@ -575,6 +812,9 @@ module ArrowFormat
       @child = child
     end
 
+    def to_s
+      "#{super}<#{child.name}: #{child.type}>"
+    end
   end
 
   class ListType < VariableSizeListType
@@ -582,9 +822,17 @@ module ArrowFormat
       "List"
     end
 
+    def offset_buffer_type
+      :s32 # TODO: big endian support
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, child)
       ListArray.new(self, size, validity_buffer, offsets_buffer, child)
     end
+
+    def to_flatbuffers
+      FB::List::Data.new
+    end
   end
 
   class LargeListType < VariableSizeListType
@@ -592,9 +840,45 @@ module ArrowFormat
       "LargeList"
     end
 
+    def offset_buffer_type
+      :s64 # TODO: big endian support
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, child)
       LargeListArray.new(self, size, validity_buffer, offsets_buffer, child)
     end
+
+    def to_flatbuffers
+      FB::LargeList::Data.new
+    end
+  end
+
+  class FixedSizeListType < Type
+    attr_reader :child
+    attr_reader :size
+    def initialize(child, size)
+      super()
+      @child = child
+      @size = size
+    end
+
+    def name
+      "FixedSizeList"
+    end
+
+    def to_s
+      "#{super}<#{child.name}: #{child.type}>(#{@size})"
+    end
+
+    def build_array(size, validity_buffer, child)
+      FixedSizeListArray.new(self, size, validity_buffer, child)
+    end
+
+    def to_flatbuffers
+      fb_type = FB::FixedSizeList::Data.new
+      fb_type.list_size = @size
+      fb_type
+    end
   end
 
   class StructType < Type
@@ -611,10 +895,21 @@ module ArrowFormat
     def build_array(size, validity_buffer, children)
       StructArray.new(self, size, validity_buffer, children)
     end
+
+    def to_s
+      fields = children.collect do |child|
+        "#{child.name}: #{child.type}"
+      end
+      "#{super}<#{fields.join(", ")}>"
+    end
+
+    def to_flatbuffers
+      FB::Struct::Data.new
+    end
   end
 
   class MapType < VariableSizeListType
-    def initialize(child)
+    def initialize(child, keys_sorted)
       if child.nullable?
         raise TypeError.new("Map entry field must not be nullable: " +
                             child.inspect)
@@ -632,22 +927,41 @@ module ArrowFormat
                             type.children[0].inspect)
       end
       super(child)
+      @keys_sorted = keys_sorted
     end
 
     def name
       "Map"
     end
 
+    def keys_sorted?
+      @keys_sorted
+    end
+
+    def offset_buffer_type
+      :s32 # TODO: big endian support
+    end
+
     def build_array(size, validity_buffer, offsets_buffer, child)
       MapArray.new(self, size, validity_buffer, offsets_buffer, child)
     end
+
+    def to_s
+      key, value, = child.type.children
+      "#{name}<#{key.type}, #{value.type}>"
+    end
+
+    def to_flatbuffers
+      FB::Map::Data.new
+    end
   end
 
   class UnionType < Type
     attr_reader :children
     attr_reader :type_ids
-    def initialize(children, type_ids)
+    def initialize(mode, children, type_ids)
       super()
+      @mode = mode
       @children = children
       @type_ids = type_ids
       @type_indexes = {}
@@ -656,19 +970,45 @@ module ArrowFormat
     def resolve_type_index(type)
       @type_indexes[type] ||= @type_ids.index(type)
     end
+
+    def to_s
+      children = @children.collect.with_index do |child, i|
+        "#{child.name}: #{child.type}=#{@type_ids[i]}"
+      end
+      "#{super}<#{children.join(", ")}>"
+    end
+
+    def to_flatbuffers
+      fb_type = FB::Union::Data.new
+      fb_type.mode = FB::UnionMode.try_convert(@mode.to_s.capitalize)
+      fb_type.type_ids = @type_ids
+      fb_type
+    end
   end
 
   class DenseUnionType < UnionType
+    def initialize(children, type_ids)
+      super(:dense, children, type_ids)
+    end
+
     def name
       "DenseUnion"
     end
 
+    def offset_buffer_type
+      :s32
+    end
+
     def build_array(size, types_buffer, offsets_buffer, children)
       DenseUnionArray.new(self, size, types_buffer, offsets_buffer, children)
     end
   end
 
   class SparseUnionType < UnionType
+    def initialize(children, type_ids)
+      super(:sparse, children, type_ids)
+    end
+
     def name
       "SparseUnion"
     end
@@ -679,26 +1019,49 @@ module ArrowFormat
   end
 
   class DictionaryType < Type
+    attr_reader :id
     attr_reader :index_type
     attr_reader :value_type
-    attr_reader :ordered
-    def initialize(index_type, value_type, ordered)
+    def initialize(id, index_type, value_type, ordered)
       super()
+      @id = id
       @index_type = index_type
       @value_type = value_type
       @ordered = ordered
     end
 
+    def ordered?
+      @ordered
+    end
+
     def name
       "Dictionary"
     end
 
-    def build_array(size, validity_buffer, indices_buffer, dictionary)
+    def build_array(size, validity_buffer, indices_buffer, dictionaries)
       DictionaryArray.new(self,
                           size,
                           validity_buffer,
                           indices_buffer,
-                          dictionary)
+                          dictionaries)
+    end
+
+    def build_fb_field(fb_field)
+      fb_dictionary_encoding = FB::DictionaryEncoding::Data.new
+      fb_dictionary_encoding.id = @id
+      fb_int = FB::Int::Data.new
+      fb_int.bit_width = @index_type.bit_width
+      fb_int.signed = @index_type.signed?
+      fb_dictionary_encoding.index_type = fb_int
+      fb_dictionary_encoding.ordered = @ordered
+      fb_dictionary_encoding.dictionary_kind = FB::DictionaryKind::DENSE_ARRAY
+      fb_field.type = @value_type.to_flatbuffers
+      fb_field.dictionary = fb_dictionary_encoding
+    end
+
+    def to_s
+      "#{super}<id=#{@id}, index=#{@index_type}, value=#{@value_type}, " +
+        "ordered=#{@ordered}>"
     end
   end
 end