red_amber 0.4.0 → 0.4.1

data/lib/red_amber/data_frame_combinable.rb

@@ -31,6 +31,8 @@ module RedAmber
     #   2       3 C
     #   3       4 D
     #
+    # @since 0.2.3
+    #
     def concatenate(*other)
       case other
       in [] | [nil] | [[]]
@@ -79,6 +81,8 @@ module RedAmber
     #   0       1       3 A        C
     #   1       2       4 B        D
     #
+    # @since 0.2.3
+    #
     def merge(*other)
       case other
       in [] | [nil] | [[]]
@@ -121,7 +125,7 @@ module RedAmber
     #   @param other [DataFrame, Arrow::Table]
     #     A DataFrame or a Table to be joined with self.
     #
-    # @!macro join_dorce_order
+    # @!macro join_force_order
     #   @param force_order [Boolean]
     #     wheather force order of the output always same.
     #     - This option is used in `:full_outer` and `:right_outer`.
@@ -217,11 +221,12 @@ module RedAmber
     # - Same as `#join` with `type: :inner`
     # - A kind of mutating join.
     #
-    # @overload inner_join(other, suffix: '.1')
+    # @overload inner_join(other, suffix: '.1', force_order: true)
     #   If `join_key` is not specified, common keys in self and other are used
     #   (natural keys). Returns joined dataframe.
     #
     #   @macro join_before
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example without key (use implicit common key)
@@ -233,10 +238,11 @@ module RedAmber
     #     0 A              1 true
     #     1 B              2 false
     #
-    # @overload inner_join(other, join_keys, suffix: '.1')
+    # @overload inner_join(other, join_keys, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_array
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example with a key
@@ -248,10 +254,11 @@ module RedAmber
     #     0 A              1 true
     #     1 B              2 false
     #
-    # @overload inner_join(other, join_key_pairs, suffix: '.1')
+    # @overload inner_join(other, join_key_pairs, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_hash
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_2
     #   @example with key pairs
@@ -263,8 +270,10 @@ module RedAmber
     #     0 A              1 true
     #     1 B              2 false
     #
-    def inner_join(other, join_keys = nil, suffix: '.1')
-      join(other, join_keys, type: :inner, suffix: suffix)
+    # @since 0.2.3
+    #
+    def inner_join(other, join_keys = nil, suffix: '.1', force_order: true)
+      join(other, join_keys, type: :inner, suffix: suffix, force_order: force_order)
     end
 
     # Join another DataFrame or Table, leaving all records.
@@ -276,7 +285,7 @@ module RedAmber
     #   (natural keys). Returns joined dataframe.
     #
     #   @macro join_before
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example without key (use implicit common key)
@@ -294,7 +303,7 @@ module RedAmber
     #
     #   @macro join_before
     #   @macro join_key_in_array
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example with a key
@@ -312,7 +321,7 @@ module RedAmber
     #
     #   @macro join_before
     #   @macro join_key_in_hash
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_2
     #   @example with key pairs
@@ -326,6 +335,8 @@ module RedAmber
     #     2 C              3 (nil)
     #     3 D          (nil) (nil)
     #
+    # @since 0.2.3
+    #
     def full_join(other, join_keys = nil, suffix: '.1', force_order: true)
       join(other, join_keys,
            type: :full_outer, suffix: suffix, force_order: force_order)
@@ -337,11 +348,12 @@ module RedAmber
     # - Same as `#join` with `type: :left_outer`
     # - A kind of mutating join.
     #
-    # @overload left_join(other, suffix: '.1')
+    # @overload left_join(other, suffix: '.1', force_order: true)
     #   If `join_key` is not specified, common keys in self and other are used
     #   (natural keys). Returns joined dataframe.
     #
     #   @macro join_before
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example without key (use implicit common key)
@@ -354,10 +366,11 @@ module RedAmber
     #     1 B              2 false
     #     2 C              3 (nil)
     #
-    # @overload left_join(other, join_keys, suffix: '.1')
+    # @overload left_join(other, join_keys, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_array
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example with a key
@@ -370,10 +383,11 @@ module RedAmber
     #     1 B              2 false
     #     2 C              3 (nil)
     #
-    # @overload left_join(other, join_key_pairs, suffix: '.1')
+    # @overload left_join(other, join_key_pairs, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_hash
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_2
     #   @example with key pairs
@@ -386,8 +400,10 @@ module RedAmber
     #     1 B              2 false
     #     2 C              3 (nil)
     #
-    def left_join(other, join_keys = nil, suffix: '.1')
-      join(other, join_keys, type: :left_outer, suffix: suffix)
+    # @since 0.2.3
+    #
+    def left_join(other, join_keys = nil, suffix: '.1', force_order: true)
+      join(other, join_keys, type: :left_outer, suffix: suffix, force_order: force_order)
     end
 
     # Join matching values from self to other.
@@ -399,7 +415,7 @@ module RedAmber
     #   (natural keys). Returns joined dataframe.
     #
     #   @macro join_before
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example without key (use implicit common key)
@@ -416,7 +432,7 @@ module RedAmber
     #
     #   @macro join_before
     #   @macro join_key_in_array
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example with a key
@@ -433,7 +449,7 @@ module RedAmber
     #
     #   @macro join_before
     #   @macro join_key_in_hash
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_2
     #   @example with key pairs
@@ -446,6 +462,8 @@ module RedAmber
     #     1 B              2 false
     #     2 D          (nil) (nil)
     #
+    # @since 0.2.3
+    #
     def right_join(other, join_keys = nil, suffix: '.1', force_order: true)
       join(
         other,
@@ -462,11 +480,12 @@ module RedAmber
     # - Same as `#join` with `type: :left_semi`
     # - A kind of filtering join.
     #
-    # @overload semi_join(other, suffix: '.1')
+    # @overload semi_join(other, suffix: '.1', force_order: true)
     #   If `join_key` is not specified, common keys in self and other are used
     #   (natural keys). Returns joined dataframe.
     #
     #   @macro join_before
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example without key (use implicit common key)
@@ -478,10 +497,11 @@ module RedAmber
     #     0 A              1
     #     1 B              2
     #
-    # @overload semi_join(other, join_keys, suffix: '.1')
+    # @overload semi_join(other, join_keys, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_array
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example with a key
@@ -493,10 +513,11 @@ module RedAmber
     #     0 A              1
     #     1 B              2
     #
-    # @overload semi_join(other, join_key_pairs, suffix: '.1')
+    # @overload semi_join(other, join_key_pairs, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_hash
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_2
     #   @example with key pairs
@@ -508,19 +529,22 @@ module RedAmber
     #     0 A              1
     #     1 B              2
     #
-    def semi_join(other, join_keys = nil, suffix: '.1')
-      join(other, join_keys, type: :left_semi, suffix: suffix)
+    # @since 0.2.3
+    #
+    def semi_join(other, join_keys = nil, suffix: '.1', force_order: true)
+      join(other, join_keys, type: :left_semi, suffix: suffix, force_order: force_order)
     end
 
     # Return records of self that do not have a match in other.
     # - Same as `#join` with `type: :left_anti`
     # - A kind of filtering join.
     #
-    # @overload anti_join(other, suffix: '.1')
+    # @overload anti_join(other, suffix: '.1', force_order: true)
     #   If `join_key` is not specified, common keys in self and other are used
     #   (natural keys). Returns joined dataframe.
     #
     #   @macro join_before
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example without key (use implicit common key)
@@ -531,10 +555,11 @@ module RedAmber
     #       <string> <uint8>
     #     0 C              3
     #
-    # @overload anti_join(other, join_keys, suffix: '.1')
+    # @overload anti_join(other, join_keys, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_array
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example with a key
@@ -545,10 +570,11 @@ module RedAmber
     #       <string> <uint8>
     #     0 C              3
     #
-    # @overload anti_join(other, join_key_pairs, suffix: '.1')
+    # @overload anti_join(other, join_key_pairs, suffix: '.1', force_order: true)
     #
     #   @macro join_before
     #   @macro join_key_in_hash
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_2
     #   @example with key pairs
@@ -559,8 +585,10 @@ module RedAmber
     #       <string> <uint8>
     #     0 C              3
     #
-    def anti_join(other, join_keys = nil, suffix: '.1')
-      join(other, join_keys, type: :left_anti, suffix: suffix)
+    # @since 0.2.3
+    #
+    def anti_join(other, join_keys = nil, suffix: '.1', force_order: true)
+      join(other, join_keys, type: :left_anti, suffix: suffix, force_order: force_order)
     end
 
     # Set operations (#intersect, #union, #difference, #set_operable?)
@@ -574,6 +602,8 @@ module RedAmber
     # @example
     #   df3.set_operable?(other3) # => true
     #
+    # @since 0.2.3
+    #
     def set_operable?(other) # rubocop:disable Naming/AccessorMethodName
       keys == other.keys.map(&:to_sym)
     end
@@ -594,6 +624,8 @@ module RedAmber
     #     <string> <uint8>
     #   0 A              1
     #
+    # @since 0.2.3
+    #
     def intersect(other)
       unless keys == other.keys.map(&:to_sym)
         raise DataFrameArgumentError, 'keys are not same with self and other'
@@ -622,6 +654,8 @@ module RedAmber
     #   3 B              4
     #   4 D              5
     #
+    # @since 0.2.3
+    #
     def union(other)
       unless keys == other.keys.map(&:to_sym)
         raise DataFrameArgumentError, 'keys are not same with self and other'
@@ -655,6 +689,8 @@ module RedAmber
     #   0 B              4
     #   1 D              5
     #
+    # @since 0.2.3
+    #
     def difference(other)
       unless keys == other.keys.map(&:to_sym)
         raise DataFrameArgumentError, 'keys are not same with self and other'
@@ -709,7 +745,7 @@ module RedAmber
     #
     #   @macro join_before
     #   @macro join_common_type
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_1
     #   @example
@@ -736,7 +772,7 @@ module RedAmber
     #   @macro join_before
     #   @macro join_key_in_array
     #   @macro join_common_type
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_3
     #   @example join keys in an Array
@@ -761,7 +797,7 @@ module RedAmber
     #   @macro join_before
     #   @macro join_key_in_hash
     #   @macro join_common_type
-    #   @macro join_dorce_order
+    #   @macro join_force_order
     #   @macro join_after
     #   @macro join_common_example_4
     #   @example without options
@@ -790,6 +826,8 @@ module RedAmber
     #     0 A        D        e
     #     1 B        E        E
     #
+    # @since 0.2.3
+    #
     def join(other, join_keys = nil, type: :inner, suffix: '.1', force_order: true)
       right_table =
         case other
@@ -804,7 +842,7 @@ module RedAmber
       type = type.to_sym
       left_index = :__LEFT_INDEX__
       right_index = :__RIGHT_INDEX__
-      if force_order && %i[full_outer right_outer].include?(type)
+      if force_order
         left_table = assign(left_index) { indices }.table
         other = DataFrame.create(other) if other.is_a?(Arrow::Table)
         right_table = other.assign(right_index) { indices }.table
@@ -820,14 +858,12 @@ module RedAmber
 
       # This is not necessary if additional procedure is contributed to Red Arrow.
       if join_keys.is_a?(Hash)
-        left_keys = join_keys[:left]
-        right_keys = join_keys[:right]
+        left_keys = ensure_keys(join_keys[:left])
+        right_keys = ensure_keys(join_keys[:right])
       else
-        left_keys = join_keys
-        right_keys = join_keys
+        left_keys = ensure_keys(join_keys)
+        right_keys = left_keys
       end
-      left_keys = Array(left_keys).map(&:to_s)
-      right_keys = Array(right_keys).map(&:to_s)
 
       case type
       when :full_outer, :left_semi, :left_anti, :right_semi, :right_anti
@@ -853,26 +889,36 @@ module RedAmber
 
       case type
       when :inner, :left_outer, :left_semi, :left_anti, :right_semi, :right_anti
-        if joined_table.keys.uniq!
-          DataFrame.create(rename_table(joined_table, n_keys, suffix))
-        else
-          DataFrame.create(joined_table)
-        end
+        dataframe =
+          if joined_table.keys.uniq!
+            DataFrame.create(rename_table(joined_table, n_keys, suffix))
+          else
+            DataFrame.create(joined_table)
+          end
+        sorter =
+          case type
+          when :inner, :left_outer
+            [left_index, right_index]
+          when :left_semi, :left_anti
+            [left_index]
+          when :right_semi, :right_anti
+            [right_index]
+          end
       when :full_outer
+        key_index_lr =
+          left_keys.map { left_table.keys.index(_1) }
+            .zip(right_keys.map { left_table.keys.size + right_table.keys.index(_1) })
         renamed_table = rename_table(joined_table, n_keys, suffix)
-        renamed_keys = renamed_table.keys
         dropper = []
-        dataframe = DataFrame.create(renamed_table).assign do |df|
-          left_keys.map do |left_key|
-            i_left_key = renamed_keys.index(left_key)
-            right_key = renamed_keys[i_left_key + table_keys.size]
-            dropper << right_key
-            [left_key.to_sym, merge_array(df[left_key].data, df[right_key].data)]
+        dataframe =
+          DataFrame.create(renamed_table).assign do |df|
+            key_index_lr.map do |l, r|
+              dropper << df.keys[r]
+              [df.keys[l], merge_array(df.vectors[l].data, df.vectors[r].data)]
+            end
           end
-        end
-        dataframe = dataframe.sort(left_index, right_index) if force_order
-
-        dataframe.drop(dropper, left_index, right_index)
+        dataframe = dataframe.drop(dropper)
+        sorter = [left_index, right_index]
       when :right_outer
         dataframe =
           if joined_table.keys.uniq!
@@ -880,20 +926,26 @@ module RedAmber
           else
             DataFrame.create(joined_table)
           end
-        if force_order
-          dataframe =
-            dataframe
-              .sort(left_index, right_index)
-              .drop(left_index, right_index)
-        end
-        dataframe.pick do
-          [right_keys, keys.map(&:to_s) - right_keys]
-        end
+        dataframe = dataframe.pick(right_keys, dataframe.keys - right_keys)
+        sorter = [left_index, right_index]
+      end
+
+      if force_order
+        dataframe
+          .sort(sorter)
+          .drop(sorter)
+      else
+        dataframe
       end
     end
 
     private
 
+    # To ensure Array of Symbols
+    def ensure_keys(keys)
+      Array(keys).map(&:to_sym)
+    end
+
     # Rename duplicate keys by suffix
     def rename_table(joined_table, n_keys, suffix)
       joined_keys = joined_table.keys
@@ -903,17 +955,9 @@ module RedAmber
       renamed_right_keys =
         other_keys.map do |key|
           if dup_keys.include?(key)
-            new_key = nil
-            loop do
-              new_key = "#{key}#{suffix}"
-              break unless joined_keys.include?(new_key)
-
-              s = suffix.succ
-              raise DataFrameArgumentError, "suffix #{suffix} is invalid" if s == suffix
-
-              suffix = s
-            end
-            new_key
+            suffixed = "#{key}#{suffix}".to_sym
+            # Find a key from suffixed.succ
+            (suffixed..).find { !joined_keys.include?(_1) }
           else
             key
           end