red-arrow 11.0.0 → 12.0.1

data/lib/arrow/table.rb

@@ -316,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)
@@ -339,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, " +
@@ -472,18 +470,22 @@ module Arrow
     #
     #     If both of `left_outputs` and `right_outputs` aren't
     #     specified, all columns in `self` and `right` are
-    #     outputted.
+    #     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
-    #     outputted.
+    #     output.
     #   @return [Arrow::Table]
     #     The joined `Arrow::Table`.
     #
     # @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.
+    #   If key(s) are not supplied, common keys in self and right are used
+    #   (natural join).
+    #
+    #   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
@@ -493,13 +495,19 @@ module Arrow
     # @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
@@ -516,8 +524,16 @@ module Arrow
     #   @macro join_common_after
     #
     # @since 7.0.0
-    def join(right, keys=nil, 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)
@@ -533,21 +549,43 @@ 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 = sink_node_options.get_reader(process_node.output_schema)
       table = reader.read_all
       share_input(table)
       table
@@ -620,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/union-array-builder.rb

@@ -0,0 +1,59 @@
+# 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 UnionArrayBuilder
+    def append_values(values, is_valids=nil)
+      if is_valids
+        is_valids.each_with_index do |is_valid, i|
+          if is_valid
+            append_value(values[i])
+          else
+            append_null
+          end
+        end
+      else
+        values.each do |value|
+          append_value(value)
+        end
+      end
+    end
+
+    alias_method :append_child_raw, :append_child
+    def append_child(builder, filed_name=nil)
+      @child_infos = nil
+      append_child_raw(builder, field_name)
+    end
+
+    private
+    def child_infos
+      @child_infos ||= create_child_infos
+    end
+
+    def create_child_infos
+      infos = {}
+      type = value_data_type
+      type.fields.zip(children, type.type_codes).each do |field, child, id|
+        infos[field.name] = {
+          builder: child,
+          id: id,
+        }
+      end
+      infos
+    end
+  end
+end

data/lib/arrow/version.rb

@@ -16,7 +16,7 @@
 # under the License.
 
 module Arrow
-  VERSION = "11.0.0"
+  VERSION = "12.0.1"
 
   module Version
     numbers, TAG = VERSION.split("-")

data/test/raw-records/test-dense-union-array.rb

@@ -76,12 +76,25 @@ module RawRecordsDenseUnionArrayTests
                            [union_array])
   end
 
+  def remove_field_names(records)
+    records.collect do |record|
+      record.collect do |column|
+        if column.nil?
+          column
+        else
+          column.values[0]
+        end
+      end
+    end
+  end
+
   def test_null
     records = [
       [{"0" => nil}],
     ]
     target = build(:null, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_boolean
@@ -90,7 +103,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:boolean, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_int8
@@ -99,7 +113,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:int8, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_uint8
@@ -108,7 +123,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:uint8, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_int16
@@ -117,7 +133,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:int16, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_uint16
@@ -126,7 +143,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:uint16, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_int32
@@ -135,7 +153,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:int32, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_uint32
@@ -144,7 +163,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:uint32, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_int64
@@ -153,7 +173,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:int64, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_uint64
@@ -162,7 +183,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:uint64, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_float
@@ -171,7 +193,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:float, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_double
@@ -180,7 +203,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:double, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_binary
@@ -189,7 +213,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:binary, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_string
@@ -198,7 +223,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:string, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_date32
@@ -207,7 +233,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:date32, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_date64
@@ -216,7 +243,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:date64, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_timestamp_second
@@ -229,7 +257,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :second,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_timestamp_milli
@@ -242,7 +271,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :milli,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_timestamp_micro
@@ -255,7 +285,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :micro,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_timestamp_nano
@@ -268,7 +299,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :nano,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_time32_second
@@ -283,7 +315,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :second,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_time32_milli
@@ -298,7 +331,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :milli,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_time64_micro
@@ -313,7 +347,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :micro,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_time64_nano
@@ -328,7 +363,8 @@ module RawRecordsDenseUnionArrayTests
                      unit: :nano,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_decimal128
@@ -342,7 +378,8 @@ module RawRecordsDenseUnionArrayTests
                      scale: 2,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_decimal256
@@ -356,7 +393,8 @@ module RawRecordsDenseUnionArrayTests
                      scale: 2,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_month_interval
@@ -365,7 +403,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:month_interval, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_day_time_interval
@@ -374,7 +413,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:day_time_interval, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_month_day_nano_interval
@@ -383,7 +423,8 @@ module RawRecordsDenseUnionArrayTests
       [{"1" => nil}],
     ]
     target = build(:month_day_nano_interval, records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_list
@@ -399,7 +440,8 @@ module RawRecordsDenseUnionArrayTests
                      },
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_struct
@@ -418,7 +460,8 @@ module RawRecordsDenseUnionArrayTests
                      ],
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_map
@@ -432,14 +475,15 @@ module RawRecordsDenseUnionArrayTests
                      item: :boolean,
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 
   def test_sparse_union
-    omit("Need to add support for SparseUnionArrayBuilder")
     records = [
       [{"0" => {"field1" => true}}],
       [{"1" => nil}],
+      [{"0" => {"field2" => 29}}],
       [{"0" => {"field2" => nil}}],
     ]
     target = build({
@@ -457,14 +501,15 @@ module RawRecordsDenseUnionArrayTests
                      type_codes: [0, 1],
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(remove_field_names(records)),
+                 target.raw_records)
   end
 
   def test_dense_union
-    omit("Need to add support for DenseUnionArrayBuilder")
     records = [
       [{"0" => {"field1" => true}}],
       [{"1" => nil}],
+      [{"0" => {"field2" => 29}}],
       [{"0" => {"field2" => nil}}],
     ]
     target = build({
@@ -482,25 +527,25 @@ module RawRecordsDenseUnionArrayTests
                      type_codes: [0, 1],
                    },
                    records)
-    assert_equal(records, target.raw_records)
+    assert_equal(remove_field_names(remove_field_names(records)),
+                 target.raw_records)
   end
 
   def test_dictionary
-    omit("Need to add support for DictionaryArrayBuilder")
     records = [
       [{"0" => "Ruby"}],
       [{"1" => nil}],
       [{"0" => "GLib"}],
     ]
-    dictionary = Arrow::StringArray.new(["GLib", "Ruby"])
     target = build({
-                                        type: :dictionary,
-                                        index_data_type: :int8,
-                                        dictionary: dictionary,
-                                        ordered: true,
-                                      },
-                                      records)
-    assert_equal(records, target.raw_records)
+                     type: :dictionary,
+                     index_data_type: :int8,
+                     value_data_type: :string,
+                     ordered: false,
+                   },
+                   records)
+    assert_equal(remove_field_names(records),
+                 target.raw_records)
   end
 end