red-arrow 8.0.0 → 24.0.0

data/ext/arrow/raw-records.cpp

@@ -84,12 +84,13 @@ namespace red_arrow {
       VISIT(UInt16)
       VISIT(UInt32)
       VISIT(UInt64)
-      // TODO
-      // VISIT(HalfFloat)
+      VISIT(HalfFloat)
       VISIT(Float)
       VISIT(Double)
       VISIT(Binary)
+      VISIT(LargeBinary)
       VISIT(String)
+      VISIT(LargeString)
       VISIT(FixedSizeBinary)
       VISIT(Date32)
       VISIT(Date64)
@@ -99,7 +100,10 @@ namespace red_arrow {
       VISIT(MonthInterval)
       VISIT(DayTimeInterval)
       VISIT(MonthDayNanoInterval)
+      VISIT(Duration)
       VISIT(List)
+      VISIT(LargeList)
+      VISIT(FixedSizeList)
       VISIT(Struct)
       VISIT(Map)
       VISIT(SparseUnion)
@@ -145,6 +149,133 @@ namespace red_arrow {
       // The number of columns.
       const int n_columns_;
     };
+
+    class RawRecordsProducer : private Converter, public arrow::ArrayVisitor {
+    public:
+      explicit RawRecordsProducer()
+        : Converter(),
+          record_(Qnil),
+          column_index_(0),
+          row_offset_(0) {
+      }
+
+      void produce(const arrow::RecordBatch& record_batch) {
+        rb::protect([&] {
+          const auto n_columns = record_batch.num_columns();
+          const auto n_rows = record_batch.num_rows();
+          for (int64_t i = 0; i < n_rows; ++i) {
+            record_ = rb_ary_new_capa(n_columns);
+            row_offset_ = i;
+            for (int i = 0; i < n_columns; ++i) {
+              const auto array = record_batch.column(i).get();
+              column_index_ = i;
+              check_status(array->Accept(this),
+                           "[record-batch][each-raw-record]");
+            }
+            rb_yield(record_);
+          }
+          return Qnil;
+        });
+      }
+
+      void produce(const arrow::Table& table) {
+        rb::protect([&] {
+          const auto n_columns = table.num_columns();
+          const auto n_rows = table.num_rows();
+          std::vector<int> chunk_indexes(n_columns);
+          std::vector<int64_t> row_offsets(n_columns);
+          for (int64_t i_row = 0; i_row < n_rows; ++i_row) {
+            record_ = rb_ary_new_capa(n_columns);
+            for (int i_column = 0; i_column < n_columns; ++i_column) {
+              column_index_ = i_column;
+              const auto chunked_array = table.column(i_column).get();
+              auto& chunk_index = chunk_indexes[i_column];
+              auto& row_offset = row_offsets[i_column];
+              auto array = chunked_array->chunk(chunk_index).get();
+              while (array->length() == row_offset) {
+                ++chunk_index;
+                row_offset = 0;
+                array = chunked_array->chunk(chunk_index).get();
+              }
+              row_offset_ = row_offset;
+              check_status(array->Accept(this),
+                           "[table][each-raw-record]");
+              ++row_offset;
+            }
+            rb_yield(record_);
+          }
+
+          return Qnil;
+        });
+      }
+
+#define VISIT(TYPE)                                                     \
+      arrow::Status Visit(const arrow::TYPE ## Array& array) override { \
+        convert(array);                                                 \
+        return arrow::Status::OK();                                     \
+      }
+
+      VISIT(Null)
+      VISIT(Boolean)
+      VISIT(Int8)
+      VISIT(Int16)
+      VISIT(Int32)
+      VISIT(Int64)
+      VISIT(UInt8)
+      VISIT(UInt16)
+      VISIT(UInt32)
+      VISIT(UInt64)
+      VISIT(HalfFloat)
+      VISIT(Float)
+      VISIT(Double)
+      VISIT(Binary)
+      VISIT(LargeBinary)
+      VISIT(String)
+      VISIT(LargeString)
+      VISIT(FixedSizeBinary)
+      VISIT(Date32)
+      VISIT(Date64)
+      VISIT(Time32)
+      VISIT(Time64)
+      VISIT(Timestamp)
+      VISIT(MonthInterval)
+      VISIT(DayTimeInterval)
+      VISIT(MonthDayNanoInterval)
+      VISIT(Duration)
+      VISIT(List)
+      VISIT(LargeList)
+      VISIT(FixedSizeList)
+      VISIT(Struct)
+      VISIT(Map)
+      VISIT(SparseUnion)
+      VISIT(DenseUnion)
+      VISIT(Dictionary)
+      VISIT(Decimal128)
+      VISIT(Decimal256)
+      // TODO
+      // VISIT(Extension)
+
+#undef VISIT
+
+    private:
+        template <typename ArrayType>
+        void convert(const ArrayType& array) {
+          auto value = Qnil;
+          if (!array.IsNull(row_offset_)) {
+            value = convert_value(array, row_offset_);
+          }
+          rb_ary_store(record_, column_index_, value);
+        }
+
+        // Destination for converted record.
+        VALUE record_;
+
+        // The current column index.
+        int column_index_;
+
+        // The current row offset.
+        int64_t row_offset_;
+    };
   }
 
   VALUE
@@ -182,4 +313,36 @@ namespace red_arrow {
 
     return records;
   }
+
+  VALUE
+  record_batch_each_raw_record(VALUE rb_record_batch) {
+    auto garrow_record_batch = GARROW_RECORD_BATCH(RVAL2GOBJ(rb_record_batch));
+    auto record_batch = garrow_record_batch_get_raw(garrow_record_batch).get();
+    RETURN_SIZED_ENUMERATOR(rb_record_batch, 0, nullptr, record_batch->num_rows());
+
+    try {
+      RawRecordsProducer producer;
+      producer.produce(*record_batch);
+    } catch (rb::State& state) {
+      state.jump();
+    }
+
+    return Qnil;
+  }
+
+  VALUE
+  table_each_raw_record(VALUE rb_table) {
+    auto garrow_table = GARROW_TABLE(RVAL2GOBJ(rb_table));
+    auto table = garrow_table_get_raw(garrow_table).get();
+    RETURN_SIZED_ENUMERATOR(rb_table, 0, nullptr, table->num_rows());
+
+    try {
+      RawRecordsProducer producer;
+      producer.produce(*table);
+    } catch (rb::State& state) {
+      state.jump();
+    }
+
+    return Qnil;
+  }
 }

data/ext/arrow/red-arrow.hpp

@@ -59,6 +59,8 @@ namespace red_arrow {
 
   VALUE record_batch_raw_records(VALUE obj);
   VALUE table_raw_records(VALUE obj);
+  VALUE record_batch_each_raw_record(VALUE obj);
+  VALUE table_each_raw_record(VALUE obj);
 
   inline VALUE time_unit_to_scale(const arrow::TimeUnit::type unit) {
     switch (unit) {

data/ext/arrow/values.cpp

@@ -65,12 +65,13 @@ namespace red_arrow {
       VISIT(UInt16)
       VISIT(UInt32)
       VISIT(UInt64)
-      // TODO
-      // VISIT(HalfFloat)
+      VISIT(HalfFloat)
       VISIT(Float)
       VISIT(Double)
       VISIT(Binary)
+      VISIT(LargeBinary)
       VISIT(String)
+      VISIT(LargeString)
       VISIT(FixedSizeBinary)
       VISIT(Date32)
       VISIT(Date64)
@@ -80,7 +81,10 @@ namespace red_arrow {
       VISIT(MonthInterval)
       VISIT(DayTimeInterval)
       VISIT(MonthDayNanoInterval)
+      VISIT(Duration)
       VISIT(List)
+      VISIT(LargeList)
+      VISIT(FixedSizeList)
       VISIT(Struct)
       VISIT(Map)
       VISIT(SparseUnion)

data/lib/arrow/array-builder.rb

@@ -33,6 +33,11 @@ module Arrow
         end
         if builder_info
           builder = builder_info[:builder]
+          if builder.nil? and builder_info[:builder_type]
+            builder = create_builder(builder_info)
+          end
+        end
+        if builder
           builder.build(values)
         else
           Arrow::StringArray.new(values)
@@ -69,14 +74,23 @@ module Arrow
             detected: true,
           }
         when Integer
-          if value < 0
+          builder_info ||= {}
+          min = builder_info[:min] || value
+          max = builder_info[:max] || value
+          min = value if value < min
+          max = value if value > max
+
+          if builder_info[:builder_type] == :int || value < 0
             {
-              builder: IntArrayBuilder.new,
-              detected: true,
+              builder_type: :int,
+              min: min,
+              max: max,
             }
           else
             {
-              builder: UIntArrayBuilder.new,
+              builder_type: :uint,
+              min: min,
+              max: max,
             }
           end
         when Time
@@ -121,28 +135,44 @@ module Arrow
             detected: true,
           }
         when BigDecimal
-          if value.to_arrow.is_a?(Decimal128)
-            {
-              builder: Decimal128ArrayBuilder.new,
-            }
-          else
+          builder_info ||= {}
+          if builder_info[:builder] or value.nan? or value.infinite?
             {
-              builder: Decimal256ArrayBuilder.new,
+              builder: StringArrayBuilder.new,
               detected: true,
             }
+          else
+            precision = [builder_info[:precision] || 0, value.precision].max
+            scale = [builder_info[:scale] || 0, value.scale].max
+            if precision <= Decimal128DataType::MAX_PRECISION
+              {
+                builder_type: :decimal128,
+                precision: precision,
+                scale: scale,
+              }
+            else
+              {
+                builder_type: :decimal256,
+                precision: precision,
+                scale: scale,
+              }
+            end
           end
         when ::Array
-          sub_builder_info = nil
+          sub_builder_info = builder_info && builder_info[:value_builder_info]
           value.each do |sub_value|
             sub_builder_info = detect_builder_info(sub_value, sub_builder_info)
             break if sub_builder_info and sub_builder_info[:detected]
           end
-          if sub_builder_info and sub_builder_info[:detected]
-            sub_value_data_type = sub_builder_info[:builder].value_data_type
+          if sub_builder_info
+            sub_builder = sub_builder_info[:builder] || create_builder(sub_builder_info)
+            return sub_builder_info unless sub_builder
+            sub_value_data_type = sub_builder.value_data_type
             field = Field.new("item", sub_value_data_type)
             {
               builder: ListArrayBuilder.new(ListDataType.new(field)),
-              detected: true,
+              value_builder_info: sub_builder_info,
+              detected: sub_builder_info[:detected],
             }
           else
             builder_info
@@ -154,6 +184,51 @@ module Arrow
           }
         end
       end
+
+      def create_builder(builder_info)
+        builder_type = builder_info[:builder_type]
+        case builder_type
+        when :decimal128
+          data_type = Decimal128DataType.new(builder_info[:precision],
+                                             builder_info[:scale])
+          Decimal128ArrayBuilder.new(data_type)
+        when :decimal256
+          data_type = Decimal256DataType.new(builder_info[:precision],
+                                             builder_info[:scale])
+          Decimal256ArrayBuilder.new(data_type)
+        when :int
+          min = builder_info[:min]
+          max = builder_info[:max]
+
+          if GLib::MININT8 <= min && max <= GLib::MAXINT8
+            Int8ArrayBuilder.new
+          elsif GLib::MININT16 <= min && max <= GLib::MAXINT16
+            Int16ArrayBuilder.new
+          elsif GLib::MININT32 <= min && max <= GLib::MAXINT32
+            Int32ArrayBuilder.new
+          elsif GLib::MININT64 <= min && max <= GLib::MAXINT64
+            Int64ArrayBuilder.new
+          else
+            StringArrayBuilder.new
+          end
+        when :uint
+          max = builder_info[:max]
+
+          if max <= GLib::MAXUINT8
+            UInt8ArrayBuilder.new
+          elsif max <= GLib::MAXUINT16
+            UInt16ArrayBuilder.new
+          elsif max <= GLib::MAXUINT32
+            UInt32ArrayBuilder.new
+          elsif max <= GLib::MAXUINT64
+            UInt64ArrayBuilder.new
+          else
+            StringArrayBuilder.new
+          end
+        else
+          nil
+        end
+      end
     end
 
     def build(values)

data/{test/test-time32-data-type.rb → lib/arrow/array-computable.rb}

@@ -15,28 +15,36 @@
 # specific language governing permissions and limitations
 # under the License.
 
-class Time32DataTypeTest < Test::Unit::TestCase
-  sub_test_case(".new") do
-    test("Arrow::TimeUnit") do
-      assert_equal("time32[ms]",
-                   Arrow::Time32DataType.new(Arrow::TimeUnit::MILLI).to_s)
+module Arrow
+  module ArrayComputable
+    def min(options: nil)
+      compute("min", options: options).value
     end
 
-    test("Symbol") do
-      assert_equal("time32[ms]",
-                   Arrow::Time32DataType.new(:milli).to_s)
+    def max(options: nil)
+      compute("max", options: options).value
     end
 
-    test("unit: Arrow::TimeUnit") do
-      data_type = Arrow::Time32DataType.new(unit: Arrow::TimeUnit::MILLI)
-      assert_equal("time32[ms]",
-                   data_type.to_s)
+    def uniq
+      unique.values
     end
 
-    test("unit: Symbol") do
-      data_type = Arrow::Time32DataType.new(unit: :milli)
-      assert_equal("time32[ms]",
-                   data_type.to_s)
+    # Finds the index of the first occurrence of a given value.
+    #
+    # @param value [Object] The value to be compared.
+    #
+    # @return [Integer] The index of the first occurrence of a given
+    #   value on found, -1 on not found.
+    #
+    # @since 12.0.0
+    def index(value)
+      value = Scalar.resolve(value, value_data_type)
+      compute("index", options: {value: value}).value
+    end
+
+    private
+    def compute(name, options: nil)
+      Function.find(name).execute([self], options).value
     end
   end
 end

data/{test/test-buffer.rb → lib/arrow/array-statistics.rb}

@@ -15,33 +15,28 @@
 # specific language governing permissions and limitations
 # under the License.
 
-class BufferTest < Test::Unit::TestCase
-  sub_test_case(".new") do
-    test("GC") do
-      data = "Hello"
-      data_id = data.object_id
-      _buffer = Arrow::Buffer.new(data)
-      data = nil
-      GC.start
-      assert_equal("Hello", ObjectSpace._id2ref(data_id))
-    end
-  end
-
-  sub_test_case("instance methods") do
-    def setup
-      @buffer = Arrow::Buffer.new("Hello")
-    end
-
-    sub_test_case("#==") do
-      test("Arrow::Buffer") do
-        assert do
-          @buffer == @buffer
+module Arrow
+  class ArrayStatistics
+    if method_defined?(:null_count_exact)
+      alias_method :null_count_raw, :null_count
+      def null_count
+        return nil unless has_null_count?
+        if null_count_exact?
+          null_count_exact
+        else
+          null_count_approximate
         end
       end
+    end
 
-      test("not Arrow::Buffer") do
-        assert do
-          not (@buffer == 29)
+    if method_defined?(:distinct_count_exact)
+      alias_method :distinct_count_raw, :distinct_count
+      def distinct_count
+        return nil unless has_distinct_count?
+        if distinct_count_exact?
+          distinct_count_exact
+        else
+          distinct_count_approximate
         end
       end
     end

data/lib/arrow/array.rb

@@ -18,8 +18,11 @@
 module Arrow
   class Array
     include Enumerable
+
+    include ArrayComputable
     include GenericFilterable
     include GenericTakeable
+    include InputReferable
 
     class << self
       def new(*args)
@@ -30,9 +33,32 @@ module Arrow
       end
 
       def builder_class
-        builder_class_name = "#{name}Builder"
-        return nil unless const_defined?(builder_class_name)
-        const_get(builder_class_name)
+        local_name = name.split("::").last
+        builder_class_name = "#{local_name}Builder"
+        return nil unless Arrow.const_defined?(builder_class_name)
+        Arrow.const_get(builder_class_name)
+      end
+
+      # @api private
+      def try_convert(value)
+        case value
+        when ::Array
+          begin
+            new(value)
+          rescue ArgumentError
+            nil
+          end
+        else
+          if value.respond_to?(:to_arrow_array)
+            begin
+              value.to_arrow_array
+            rescue RangeError
+              nil
+            end
+          else
+            nil
+          end
+        end
       end
     end
 
@@ -67,6 +93,8 @@ module Arrow
       equal_options(other, options)
     end
 
+    alias_method :size, :length
+
     def each
       return to_enum(__method__) unless block_given?
 
@@ -87,6 +115,14 @@ module Arrow
       self
     end
 
+    def to_arrow_array
+      self
+    end
+
+    def to_arrow_chunked_array
+      ChunkedArray.new([self])
+    end
+
     alias_method :value_data_type_raw, :value_data_type
     def value_data_type
       @value_data_type ||= value_data_type_raw
@@ -217,7 +253,7 @@ module Arrow
             "[array][resolve] need to implement " +
             "a feature that building #{value_data_type} array " +
             "from raw Ruby Array"
-          raise NotImpelemented, message
+          raise NotImplementedError, message
         end
         other_array
       elsif other_array.respond_to?(:value_data_type)

data/lib/arrow/chunked-array.rb

@@ -18,8 +18,35 @@
 module Arrow
   class ChunkedArray
     include Enumerable
+
+    include ArrayComputable
     include GenericFilterable
     include GenericTakeable
+    include InputReferable
+
+    def freeze
+      unless frozen?
+        # Ensure caching
+        chunks
+      end
+      super
+    end
+
+    def to_arrow
+      self
+    end
+
+    def to_arrow_array
+      if n_chunks.zero?
+        value_data_type.build_array([])
+      else
+        combine
+      end
+    end
+
+    def to_arrow_chunked_array
+      self
+    end
 
     alias_method :size, :n_rows
     unless method_defined?(:length)
@@ -28,7 +55,16 @@ module Arrow
 
     alias_method :chunks_raw, :chunks
     def chunks
-      @chunks ||= chunks_raw
+      @chunks ||= chunks_raw.tap do |_chunks|
+        _chunks.each do |chunk|
+          share_input(chunk)
+        end
+      end
+    end
+
+    alias_method :get_chunk_raw, :get_chunk
+    def get_chunk(i)
+      chunks[i]
     end
 
     def null?(i)
@@ -87,5 +123,24 @@ module Arrow
         first_chunk.class.new(to_a)
       end
     end
+
+    def count(options: nil)
+      compute("count", options: options).value
+    end
+
+    def sum(options: nil)
+      compute("sum", options: options).value
+    end
+
+    def unique
+      compute("unique")
+    end
+
+    def cast(target_data_type, options: nil)
+      casted_chunks = chunks.collect do |chunk|
+        chunk.cast(target_data_type, options)
+      end
+      self.class.new(casted_chunks)
+    end
   end
 end

data/lib/arrow/column-containable.rb

@@ -143,5 +143,14 @@ module Arrow
         find_column(selector)
       end
     end
+
+    # Return column names in this object.
+    #
+    # @return [::Array<String>] column names.
+    #
+    # @since 11.0.0
+    def column_names
+      @column_names ||= columns.collect(&:name)
+    end
   end
 end