red-arrow 8.0.0 → 24.0.0

data/lib/arrow/table.rb

@@ -22,6 +22,7 @@ module Arrow
     include ColumnContainable
     include GenericFilterable
     include GenericTakeable
+    include InputReferable
     include RecordContainable
 
     class << self
@@ -126,7 +127,7 @@ module Arrow
     #     You can also specify schema as primitive Ruby objects.
     #     See {Arrow::Schema#initialize} for details.
     #
-    #   @param arrays [::Array<Arrow::RecordBatch>] The data of the table.
+    #   @param record_batches [::Array<Arrow::RecordBatch>] The data of the table.
     #
     #   @example Create a table from schema and record batches
     #     count_field = Arrow::Field.new("count", :uint32)
@@ -144,7 +145,7 @@ module Arrow
     #     You can also specify schema as primitive Ruby objects.
     #     See {Arrow::Schema#initialize} for details.
     #
-    #   @param arrays [::Array<::Array>] The data of the table as primitive
+    #   @param raw_records [::Array<::Array>] The data of the table as primitive
     #     Ruby objects.
     #
     #   @example Create a table from schema and raw records
@@ -188,6 +189,7 @@ module Arrow
 
       reader = TableBatchReader.new(self)
       while record_batch = reader.read_next
+        share_input(record_batch)
         yield(record_batch)
       end
     end
@@ -314,8 +316,6 @@ module Arrow
         end
       end
 
-      filter_options = Arrow::FilterOptions.new
-      filter_options.null_selection_behavior = :emit_null
       sliced_tables = []
       slicers.each do |slicer|
         slicer = slicer.evaluate if slicer.respond_to?(:evaluate)
@@ -337,7 +337,7 @@ module Arrow
           to += n_rows if to < 0
           sliced_tables << slice_by_range(from, to)
         when ::Array, BooleanArray, ChunkedArray
-          sliced_tables << filter(slicer, filter_options)
+          sliced_tables << filter(slicer)
         else
           message = "slicer must be Integer, Range, (from, to), " +
             "Arrow::ChunkedArray of Arrow::BooleanArray, " +
@@ -346,10 +346,12 @@ module Arrow
         end
       end
       if sliced_tables.size > 1
-        sliced_tables[0].concatenate(sliced_tables[1..-1])
+        sliced_table = sliced_tables[0].concatenate(sliced_tables[1..-1])
       else
-        sliced_tables[0]
+        sliced_table = sliced_tables[0]
       end
+      share_input(sliced_table)
+      sliced_table
     end
 
     # TODO
@@ -401,7 +403,9 @@ module Arrow
         new_fields << new_column[:field]
         new_arrays << new_column[:data]
       end
-      self.class.new(new_fields, new_arrays)
+      table = self.class.new(new_fields, new_arrays)
+      share_input(table)
+      table
     end
 
     alias_method :remove_column_raw, :remove_column
@@ -423,7 +427,9 @@ module Arrow
           raise IndexError.new(message)
         end
       end
-      remove_column_raw(index)
+      table = remove_column_raw(index)
+      share_input(table)
+      table
     end
 
     # Experimental
@@ -445,43 +451,69 @@ module Arrow
       packed_arrays = columns.collect do |column|
         column.data.pack
       end
-      self.class.new(schema, packed_arrays)
+      table = self.class.new(schema, packed_arrays)
+      share_input(table)
+      table
     end
 
-    # @overload join(right, key, type: :inner, left_outputs: nil, right_outputs: nil)
-    #   @!macro join_common_before
-    #     @param right [Arrow::Table] The right table.
+    # Join another Table by matching with keys.
+    #
+    # @!macro join_common_before
+    #   @param right [Arrow::Table] The right table.
+    #
+    #   Join columns with `right` on join key columns.
     #
-    #     Join columns with `right` on join key columns.
+    # @!macro join_common_after
+    #   @param type [Arrow::JoinType] How to join.
+    #   @param left_outputs [::Array<String, Symbol>] Output columns in
+    #     `self`.
     #
-    #   @!macro join_common_after
-    #     @param type [Arrow::JoinType] How to join.
-    #     @param left_outputs [::Array<String, Symbol>] Output columns in
-    #       `self`.
+    #     If both of `left_outputs` and `right_outputs` aren't
+    #     specified, all columns in `self` and `right` are
+    #     output.
+    #   @param right_outputs [::Array<String, Symbol>] Output columns in
+    #     `right`.
+    #
+    #     If both of `left_outputs` and `right_outputs` aren't
+    #     specified, all columns in `self` and `right` are
+    #     output.
+    #   @return [Arrow::Table]
+    #     The joined `Arrow::Table`.
     #
-    #       If both of `left_outputs` and `right_outputs` aren't
-    #       specified, all columns in `self` and `right` are
-    #       outputted.
-    #     @param right_outputs [::Array<String, Symbol>] Output columns in
-    #       `right`.
+    # @overload join(right, type: :inner, left_outputs: nil, right_outputs: nil)
+    #   If key(s) are not supplied, common keys in self and right are used
+    #   (natural join).
     #
-    #       If both of `left_outputs` and `right_outputs` aren't
-    #       specified, all columns in `self` and `right` are
-    #       outputted.
-    #     @return [Arrow::Table]
-    #       The joined `Arrow::Table`.
+    #   Column used as keys are merged and remain in left side
+    #   when both of `left_outputs` and `right_outputs` are `nil`.
+    #
+    #   @macro join_common_before
+    #   @macro join_common_after
+    #
+    # @since 11.0.0
+    #
+    # @overload join(right, key, type: :inner, left_outputs: nil, right_outputs: nil)
+    #   Join right by a key.
+    #
+    #   Column used as keys are merged and remain in left side
+    #   when both of `left_outputs` and `right_outputs` are `nil`.
     #
     #   @macro join_common_before
     #   @param key [String, Symbol] A join key.
     #   @macro join_common_after
     #
-    # @overload join(right, keys, type: :inner, left_outputs: nil, right_outputs: nil)
+    # @overload join(right, keys, type: :inner, left_suffix: "", right_suffix: "",
+    #                left_outputs: nil, right_outputs: nil)
+    #   Join right by keys.
+    #
+    #   Column name can be renamed by appending `left_suffix` or `right_suffix`.
     #
     #   @macro join_common_before
     #   @param keys [::Array<String, Symbol>] Join keys.
     #   @macro join_common_after
     #
     # @overload join(right, keys, type: :inner, left_outputs: nil, right_outputs: nil)
+    #   Join right by a key or keys mapped by a hash.
     #
     #   @macro join_common_before
     #   @param keys [Hash] Specify join keys in `self` and `right` separately.
@@ -492,7 +524,16 @@ module Arrow
     #   @macro join_common_after
     #
     # @since 7.0.0
-    def join(right, keys, type: :inner, left_outputs: nil, right_outputs: nil)
+    def join(right,
+             keys=nil,
+             type: :inner,
+             left_suffix: "",
+             right_suffix: "",
+             left_outputs: nil,
+             right_outputs: nil)
+      is_natural_join = keys.nil?
+      keys ||= (column_names & right.column_names)
+      type = JoinType.try_convert(type) || type
       plan = ExecutePlan.new
       left_node = plan.build_source_node(self)
       right_node = plan.build_source_node(right)
@@ -508,22 +549,46 @@ module Arrow
       hash_join_node_options = HashJoinNodeOptions.new(type,
                                                        left_keys,
                                                        right_keys)
+      use_manual_outputs = false
       unless left_outputs.nil?
         hash_join_node_options.left_outputs = left_outputs
+        use_manual_outputs = true
       end
       unless right_outputs.nil?
         hash_join_node_options.right_outputs = right_outputs
+        use_manual_outputs = true
       end
       hash_join_node = plan.build_hash_join_node(left_node,
                                                  right_node,
                                                  hash_join_node_options)
+      type_nick = type.nick
+      is_filter_join = (type_nick.end_with?("-semi") or
+                        type_nick.end_with?("-anti"))
+      if use_manual_outputs or is_filter_join
+        process_node = hash_join_node
+      elsif is_natural_join
+        process_node = join_merge_keys(plan, hash_join_node, right, keys)
+      elsif keys.is_a?(String) or keys.is_a?(Symbol)
+        process_node = join_merge_keys(plan, hash_join_node, right, [keys.to_s])
+      elsif !keys.is_a?(Hash) and (left_suffix != "" or right_suffix != "")
+        process_node = join_rename_keys(plan,
+                                        hash_join_node,
+                                        right,
+                                        keys,
+                                        left_suffix,
+                                        right_suffix)
+      else
+        process_node = hash_join_node
+      end
       sink_node_options = SinkNodeOptions.new
-      plan.build_sink_node(hash_join_node, sink_node_options)
+      plan.build_sink_node(process_node, sink_node_options)
       plan.validate
       plan.start
       plan.wait
-      reader = sink_node_options.get_reader(hash_join_node.output_schema)
-      reader.read_all
+      reader = sink_node_options.get_reader(process_node.output_schema)
+      table = reader.read_all
+      share_input(table)
+      table
     end
 
     alias_method :to_s_raw, :to_s
@@ -593,5 +658,87 @@ module Arrow
         raise ArgumentError, message
       end
     end
+
+    def join_merge_keys(plan, input_node, right, keys)
+      expressions = []
+      names = []
+      normalized_keys = {}
+      keys.each do |key|
+        normalized_keys[key.to_s] = true
+      end
+      key_to_outputs = {}
+      outputs = []
+      left_n_column_names = column_names.size
+      column_names.each_with_index do |name, i|
+        is_key = normalized_keys.include?(name)
+        output = {is_key: is_key, name: name, index: i, direction: :left}
+        outputs << output
+        key_to_outputs[name] = {left: output} if is_key
+      end
+      right.column_names.each_with_index do |name, i|
+        index = left_n_column_names + i
+        is_key = normalized_keys.include?(name)
+        output = {is_key: is_key, name: name, index: index, direction: :right}
+        outputs << output
+        key_to_outputs[name][:right] = output if is_key
+      end
+
+      outputs.each do |output|
+        if output[:is_key]
+          next if output[:direction] == :right
+          left_output = key_to_outputs[output[:name]][:left]
+          right_output = key_to_outputs[output[:name]][:right]
+          left_field = FieldExpression.new("[#{left_output[:index]}]")
+          right_field = FieldExpression.new("[#{right_output[:index]}]")
+          is_left_null = CallExpression.new("is_null", [left_field])
+          merge_column = CallExpression.new("if_else",
+                                            [
+                                              is_left_null,
+                                              right_field,
+                                              left_field,
+                                            ])
+          expressions << merge_column
+        else
+          expressions << FieldExpression.new("[#{output[:index]}]")
+        end
+        names << output[:name]
+      end
+      project_node_options = ProjectNodeOptions.new(expressions, names)
+      plan.build_project_node(input_node, project_node_options)
+    end
+
+    def join_rename_keys(plan,
+                         input_node,
+                         right,
+                         keys,
+                         left_suffix,
+                         right_suffix)
+      expressions = []
+      names = []
+      normalized_keys = {}
+      keys.each do |key|
+        normalized_keys[key.to_s] = true
+      end
+      left_n_column_names = column_names.size
+      column_names.each_with_index do |name, i|
+        expressions << FieldExpression.new("[#{i}]")
+        if normalized_keys.include?(name)
+          names << "#{name}#{left_suffix}"
+        else
+          names << name
+        end
+      end
+      right.column_names.each_with_index do |name, i|
+        index = left_n_column_names + i
+        expressions << FieldExpression.new("[#{index}]")
+        if normalized_keys.include?(name)
+          names << "#{name}#{right_suffix}"
+        else
+          names << name
+        end
+      end
+      project_node_options = ProjectNodeOptions.new(expressions, names)
+      plan.build_project_node(input_node, project_node_options)
+    end
   end
 end

data/lib/arrow/tensor.rb

@@ -15,10 +15,154 @@
 # specific language governing permissions and limitations
 # under the License.
 
+require_relative "raw-tensor-converter"
+
 module Arrow
   class Tensor
+    alias_method :initialize_raw, :initialize
+    # Creates a new {Arrow::Tensor}.
+    #
+    # @overload initialize(raw_tensor, data_type: nil, shape: nil, dimension_names: nil)
+    #
+    #   @param raw_tensor [::Array<Numeric>] The tensor represented as a
+    #     raw `Array` (not `Arrow::Array`) and `Numeric`s. You can
+    #     pass a nested `Array` for a multi-dimensional tensor.
+    #
+    #   @param data_type [Arrow::DataType, String, Symbol, ::Array<String>,
+    #     ::Array<Symbol>, Hash, nil] The element data type of the tensor.
+    #
+    #     If you specify `nil`, data type is guessed from `raw_tensor`.
+    #
+    #     See {Arrow::DataType.resolve} for how to specify data type.
+    #
+    #   @param shape [::Array<Integer>, nil] The array of dimension sizes.
+    #
+    #     If you specify `nil`, shape is guessed from `raw_tensor`.
+    #
+    #   @param dimension_names [::Array<String>, ::Array<Symbol>, nil]
+    #     The array of the dimension names.
+    #
+    #     If you specify `nil`, all dimensions have empty names.
+    #
+    #   @example Create a tensor from Ruby's Array
+    #     raw_tensor = [
+    #       [
+    #         [1, 2, 3, 4],
+    #         [5, 6, 7, 8],
+    #       ],
+    #       [
+    #         [9, 10, 11, 12],
+    #         [13, 14, 15, 16],
+    #       ],
+    #       [
+    #         [17, 18, 19, 20],
+    #         [21, 22, 23, 24],
+    #       ],
+    #     ]
+    #     Arrow::Tensor.new(raw_tensor)
+    #
+    #   @since 10.0.0
+    #
+    # @overload initialize(data_type, data, shape, strides, dimension_names)
+    #
+    #   @param data_type [Arrow::DataType, String, Symbol, ::Array<String>,
+    #     ::Array<Symbol>, Hash] The element data type of the tensor.
+    #
+    #     See {Arrow::DataType.resolve} how to specify data type.
+    #
+    #   @param data [Arrow::Buffer, String] The data of the tensor.
+    #
+    #   @param shape [::Array<Integer>] The array of dimension sizes.
+    #
+    #   @param strides [::Array<Integer>, nil] The array of strides which
+    #     is the number of bytes between two adjacent elements in each
+    #     dimension.
+    #
+    #     If you specify `nil` or an empty `Array`, strides are
+    #     guessed from `data_type` and `data`.
+    #
+    #   @param dimension_names [::Array<String>, ::Array<Symbol>, nil]
+    #     The array of the dimension names.
+    #
+    #     If you specify `nil`, all dimensions doesn't have their names.
+    #
+    #   @example Create a table from Arrow::Buffer
+    #     raw_data = [
+    #       1, 2,
+    #       3, 4,
+    #
+    #       5, 6,
+    #       7, 8,
+    #
+    #       9, 10,
+    #       11, 12,
+    #     ]
+    #     data = Arrow::Buffer.new(raw_data.pack("c*").freeze)
+    #     shape = [3, 2, 2]
+    #     strides = []
+    #     names = ["a", "b", "c"]
+    #     Arrow::Tensor.new(:int8, data, shape, strides, names)
+    def initialize(*args,
+                   data_type: nil,
+                   data: nil,
+                   shape: nil,
+                   strides: nil,
+                   dimension_names: nil)
+      n_args = args.size
+      case n_args
+      when 1
+        converter = RawTensorConverter.new(args[0],
+                                           data_type: data_type,
+                                           shape: shape,
+                                           strides: strides,
+                                           dimension_names: dimension_names)
+        data_type = converter.data_type
+        data = converter.data
+        shape = converter.shape
+        strides = converter.strides
+        dimension_names = converter.dimension_names
+      when 0, 2..5
+        data_type = args[0] || data_type
+        data = args[1] || data
+        shape = args[2] || shape
+        strides = args[3] || strides
+        dimension_names = args[4] || dimension_names
+        if data_type.nil?
+          raise ArgumentError, "data_type: is missing: #{data.inspect}"
+        end
+      else
+        message = "wrong number of arguments (given #{n_args}, expected 0..5)"
+        raise ArgumentError, message
+      end
+      initialize_raw(DataType.resolve(data_type),
+                     data,
+                     shape,
+                     strides,
+                     dimension_names)
+    end
+
+    def dimension_names
+      n_dimensions.times.collect do |i|
+        get_dimension_name(i)
+      end
+    end
+
     def to_arrow
       self
     end
+
+    def to_arrow_array
+      if n_dimensions != 1
+        raise RangeError, "must be 1 dimensional tensor: #{shape.inspect}"
+      end
+      value_data_type.array_class.new(size,
+                                      buffer,
+                                      nil,
+                                      0)
+    end
+
+    def to_arrow_chunked_array
+      ChunkedArray.new([to_arrow_array])
+    end
   end
 end

data/lib/arrow/time-unit.rb

@@ -0,0 +1,31 @@
+# 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
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+module Arrow
+  class TimeUnit
+    class << self
+      # @api private
+      def try_convert(value)
+        if value.is_a?(Hash) and value.size == 1 and value[:unit]
+          super(value[:unit])
+        else
+          super
+        end
+      end
+    end
+  end
+end

data/lib/arrow/time32-array-builder.rb

@@ -18,24 +18,12 @@
 module Arrow
   class Time32ArrayBuilder
     class << self
-      def build(unit_or_data_type, values)
-        builder = new(unit_or_data_type)
+      def build(data_type, values)
+        builder = new(data_type)
         builder.build(values)
       end
     end
 
-    alias_method :initialize_raw, :initialize
-    def initialize(unit_or_data_type)
-      case unit_or_data_type
-      when DataType
-        data_type = unit_or_data_type
-      else
-        unit = unit_or_data_type
-        data_type = Time32DataType.new(unit)
-      end
-      initialize_raw(data_type)
-    end
-
     def unit
       @unit ||= value_data_type.unit
     end

data/lib/arrow/time32-data-type.rb

@@ -17,45 +17,16 @@
 
 module Arrow
   class Time32DataType
-    alias_method :initialize_raw, :initialize
-    private :initialize_raw
-
-    # Creates a new {Arrow::Time32DataType}.
-    #
-    # @overload initialize(unit)
-    #
-    #   @param unit [Arrow::TimeUnit, Symbol] The unit of the
-    #     time32 data type.
-    #
-    #     The unit must be second or millisecond.
-    #
-    #   @example Create a time32 data type with Arrow::TimeUnit
-    #     Arrow::Time32DataType.new(Arrow::TimeUnit::MILLI)
-    #
-    #   @example Create a time32 data type with Symbol
-    #     Arrow::Time32DataType.new(:milli)
-    #
-    # @overload initialize(description)
-    #
-    #   @param description [Hash] The description of the time32 data
-    #     type. It must have `:unit` value.
-    #
-    #   @option description [Arrow::TimeUnit, Symbol] :unit The unit of
-    #     the time32 data type.
-    #
-    #     The unit must be second or millisecond.
-    #
-    #   @example Create a time32 data type with Arrow::TimeUnit
-    #     Arrow::Time32DataType.new(unit: Arrow::TimeUnit::MILLI)
-    #
-    #   @example Create a time32 data type with Symbol
-    #     Arrow::Time32DataType.new(unit: :milli)
-    def initialize(unit)
-      if unit.is_a?(Hash)
-        description = unit
-        unit = description[:unit]
+    class << self
+      # @api private
+      def try_convert(value)
+        case value
+        when Symbol, Arrow::TimeUnit
+          new(value)
+        else
+          super
+        end
       end
-      initialize_raw(unit)
     end
   end
 end

data/lib/arrow/time64-array-builder.rb

@@ -18,24 +18,12 @@
 module Arrow
   class Time64ArrayBuilder
     class << self
-      def build(unit_or_data_type, values)
-        builder = new(unit_or_data_type)
+      def build(data_type, values)
+        builder = new(data_type)
         builder.build(values)
       end
     end
 
-    alias_method :initialize_raw, :initialize
-    def initialize(unit_or_data_type)
-      case unit_or_data_type
-      when DataType
-        data_type = unit_or_data_type
-      else
-        unit = unit_or_data_type
-        data_type = Time64DataType.new(unit)
-      end
-      initialize_raw(data_type)
-    end
-
     def unit
       @unit ||= value_data_type.unit
     end

data/lib/arrow/time64-data-type.rb

@@ -17,45 +17,16 @@
 
 module Arrow
   class Time64DataType
-    alias_method :initialize_raw, :initialize
-    private :initialize_raw
-
-    # Creates a new {Arrow::Time64DataType}.
-    #
-    # @overload initialize(unit)
-    #
-    #   @param unit [Arrow::TimeUnit, Symbol] The unit of the
-    #     time64 data type.
-    #
-    #     The unit must be microsecond or nanosecond.
-    #
-    #   @example Create a time64 data type with Arrow::TimeUnit
-    #     Arrow::Time64DataType.new(Arrow::TimeUnit::NANO)
-    #
-    #   @example Create a time64 data type with Symbol
-    #     Arrow::Time64DataType.new(:nano)
-    #
-    # @overload initialize(description)
-    #
-    #   @param description [Hash] The description of the time64 data
-    #     type. It must have `:unit` value.
-    #
-    #   @option description [Arrow::TimeUnit, Symbol] :unit The unit of
-    #     the time64 data type.
-    #
-    #     The unit must be microsecond or nanosecond.
-    #
-    #   @example Create a time64 data type with Arrow::TimeUnit
-    #     Arrow::Time64DataType.new(unit: Arrow::TimeUnit::NANO)
-    #
-    #   @example Create a time64 data type with Symbol
-    #     Arrow::Time64DataType.new(unit: :nano)
-    def initialize(unit)
-      if unit.is_a?(Hash)
-        description = unit
-        unit = description[:unit]
+    class << self
+      # @api private
+      def try_convert(value)
+        case value
+        when Symbol, Arrow::TimeUnit
+          new(value)
+        else
+          super
+        end
       end
-      initialize_raw(unit)
     end
   end
 end