red_amber 0.4.1 → 0.5.0

data/lib/red_amber/data_frame_combinable.rb

@@ -221,6 +221,11 @@ module RedAmber
     # - Same as `#join` with `type: :inner`
     # - A kind of mutating join.
     #
+    # @note the order of joined results will be preserved by default.
+    #   This is enabled by appending index column to sort after joining but
+    #   it will cause some performance degradation. If you don't matter
+    #   the order of the result, set `force_order` option to `false`.
+    #
     # @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.
@@ -280,6 +285,11 @@ module RedAmber
     # - Same as `#join` with `type: :full_outer`
     # - A kind of mutating join.
     #
+    # @note the order of joined results will be preserved by default.
+    #   This is enabled by appending index column to sort after joining but
+    #   it will cause some performance degradation. If you don't matter
+    #   the order of the result, set `force_order` option to `false`.
+    #
     # @overload full_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.
@@ -348,6 +358,11 @@ module RedAmber
     # - Same as `#join` with `type: :left_outer`
     # - A kind of mutating join.
     #
+    # @note the order of joined results will be preserved by default.
+    #   This is enabled by appending index column to sort after joining but
+    #   it will cause some performance degradation. If you don't matter
+    #   the order of the result, set `force_order` option to `false`.
+    #
     # @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.
@@ -410,6 +425,11 @@ module RedAmber
     # - Same as `#join` with `type: :right_outer`
     # - A kind of mutating join.
     #
+    # @note the order of joined results will be preserved by default.
+    #   This is enabled by appending index column to sort after joining but
+    #   it will cause some performance degradation. If you don't matter
+    #   the order of the result, set `force_order` option to `false`.
+    #
     # @overload right_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.
@@ -422,11 +442,11 @@ module RedAmber
     #     df.right_join(other)
     #
     #     # =>
-    #       KEY           X1 X2
-    #       <string> <uint8> <boolean>
-    #     0 A              1 true
-    #     1 B              2 false
-    #     2 D          (nil) (nil)
+    #            X1 KEY      X2
+    #       <uint8> <string> <boolean>
+    #     0       1 A        true
+    #     1       2 B        false
+    #     2   (nil) D        (nil)
     #
     # @overload right_join(other, join_keys, suffix: '.1', force_order: true)
     #
@@ -439,11 +459,11 @@ module RedAmber
     #     df.right_join(other, :KEY)
     #
     #     # =>
-    #       KEY           X1 X2
-    #       <string> <uint8> <boolean>
-    #     0 A              1 true
-    #     1 B              2 false
-    #     2 D          (nil) (nil)
+    #            X1 KEY      X2
+    #       <uint8> <string> <boolean>
+    #     0       1 A        true
+    #     1       2 B        false
+    #     2   (nil) D        (nil)
     #
     # @overload right_join(other, join_key_pairs, suffix: '.1', force_order: true)
     #
@@ -456,11 +476,11 @@ module RedAmber
     #     df2.right_join(other2, { left: :KEY1, right: :KEY2 })
     #
     #     # =>
-    #       KEY1          X1 X2
-    #       <string> <uint8> <boolean>
-    #     0 A              1 true
-    #     1 B              2 false
-    #     2 D          (nil) (nil)
+    #             X1 KEY2     X2
+    #       <uint8> >string> <boolean>
+    #     0        1 A        true
+    #     1        2 B        false
+    #     2    (nil) D        (nil)
     #
     # @since 0.2.3
     #
@@ -480,6 +500,11 @@ module RedAmber
     # - Same as `#join` with `type: :left_semi`
     # - A kind of filtering join.
     #
+    # @note the order of joined results will be preserved by default.
+    #   This is enabled by appending index column to sort after joining but
+    #   it will cause some performance degradation. If you don't matter
+    #   the order of the result, set `force_order` option to `false`.
+    #
     # @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.
@@ -539,6 +564,11 @@ module RedAmber
     # - Same as `#join` with `type: :left_anti`
     # - A kind of filtering join.
     #
+    # @note the order of joined results will be preserved by default.
+    #   This is enabled by appending index column to sort after joining but
+    #   it will cause some performance degradation. If you don't matter
+    #   the order of the result, set `force_order` option to `false`.
+    #
     # @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.
@@ -661,7 +691,7 @@ module RedAmber
         raise DataFrameArgumentError, 'keys are not same with self and other'
       end
 
-      join(other, keys, type: :full_outer)
+      join(other, keys, type: :full_outer, force_order: true)
     end
 
     # Select records appearing in self but not in other.
@@ -733,12 +763,12 @@ module RedAmber
     #     1 B        E
     #     2 C        F
 
-    # @note the order of joined results will be preserved by default.
-    #   This is enabled by appending index column to sort after joining but
-    #   it will cause some performance degradation. If you don't matter
-    #   the order of the result, set `force_order` option to `false`.
+    # @note the order of joined results may not be preserved by default.
+    #   if you prefer to preserve the order of the result, set `force_order` option
+    #   to `true`. This is enabled by appending index column to sort after joining
+    #   so it will cause some performance degradation.
     #
-    # @overload join(other, type: :inner, suffix: '.1', force_order: true)
+    # @overload join(other, type: :inner, suffix: '.1', force_order: false)
     #
     #   If `join_key` is not specified, common keys in self and other are used
     #   (natural keys). Returns joined dataframe.
@@ -767,7 +797,7 @@ module RedAmber
     #     2 C              3 (nil)
     #     3 D          (nil) (nil)
     #
-    # @overload join(other, join_keys, type: :inner, suffix: '.1', force_order: true)
+    # @overload join(other, join_keys, type: :inner, suffix: '.1', force_order: false)
     #
     #   @macro join_before
     #   @macro join_key_in_array
@@ -792,7 +822,8 @@ module RedAmber
     #     0 A              1       1
     #     1 B              2       4
     #
-    # @overload join(other, join_key_pairs, type: :inner, suffix: '.1', force_order: true)
+    # @overload join(
+    #   other, join_key_pairs, type: :inner, suffix: '.1', force_order: false)
     #
     #   @macro join_before
     #   @macro join_key_in_hash
@@ -828,7 +859,8 @@ module RedAmber
     #
     # @since 0.2.3
     #
-    def join(other, join_keys = nil, type: :inner, suffix: '.1', force_order: true)
+    def join(other, join_keys = nil, type: :inner, suffix: '.1', force_order: false)
+      left_table = table
       right_table =
         case other
         when DataFrame
@@ -839,24 +871,26 @@ module RedAmber
           raise DataFrameArgumentError, 'other must be a DataFrame or an Arrow::Table'
         end
 
-      type = type.to_sym
-      left_index = :__LEFT_INDEX__
-      right_index = :__RIGHT_INDEX__
       if force_order
+        left_index = :__LEFT_INDEX__
+        right_index = :__RIGHT_INDEX__
         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
-      else
-        left_table = table
       end
 
-      table_keys = left_table.keys
-      other_keys = right_table.keys
-
+      left_table_keys = ensure_keys(left_table.keys)
+      right_table_keys = ensure_keys(right_table.keys)
       # natural keys (implicit common keys)
-      join_keys ||= table_keys.intersection(other_keys)
+      join_keys ||= left_table_keys.intersection(right_table_keys)
+
+      type = Arrow::JoinType.try_convert(type) || type
+      type_nick = type.nick
+
+      plan = Arrow::ExecutePlan.new
+      left_node = plan.build_source_node(left_table)
+      right_node = plan.build_source_node(right_table)
 
-      # This is not necessary if additional procedure is contributed to Red Arrow.
       if join_keys.is_a?(Hash)
         left_keys = ensure_keys(join_keys[:left])
         right_keys = ensure_keys(join_keys[:right])
@@ -865,116 +899,110 @@ module RedAmber
         right_keys = left_keys
       end
 
-      case type
-      when :full_outer, :left_semi, :left_anti, :right_semi, :right_anti
-        left_outputs = nil
-        right_outputs = nil
-      when :inner, :left_outer
-        left_outputs = table_keys
-        right_outputs = other_keys - right_keys
-      when :right_outer
-        left_outputs = table_keys - left_keys
-        right_outputs = other_keys
+      context =
+        [type_nick, left_table_keys, right_table_keys, left_keys, right_keys, suffix]
+
+      hash_join_node_options = Arrow::HashJoinNodeOptions.new(type, left_keys, right_keys)
+      case type_nick
+      when 'inner', 'left-outer'
+        hash_join_node_options.left_outputs = left_table_keys
+        hash_join_node_options.right_outputs = right_table_keys - right_keys
+      when 'right-outer'
+        hash_join_node_options.left_outputs = left_table_keys - left_keys
+        hash_join_node_options.right_outputs = right_table_keys
       end
 
-      # Should we rescue errors in Arrow::Table#join for usability ?
-      joined_table =
-        left_table.join(
-          right_table,
-          join_keys,
-          type: type,
-          left_outputs: left_outputs,
-          right_outputs: right_outputs
-        )
-
-      case type
-      when :inner, :left_outer, :left_semi, :left_anti, :right_semi, :right_anti
-        dataframe =
-          if joined_table.keys.uniq!
-            DataFrame.create(rename_table(joined_table, n_keys, suffix))
-          else
-            DataFrame.create(joined_table)
-          end
+      hash_join_node =
+        plan.build_hash_join_node(left_node, right_node, hash_join_node_options)
+      merge_node = merge_keys(plan, hash_join_node, context)
+      rename_node = rename_keys(plan, merge_node, context)
+      joined_table = sink_and_start_plan(plan, rename_node)
+
+      df = DataFrame.create(joined_table)
+      if force_order
         sorter =
-          case type
-          when :inner, :left_outer
-            [left_index, right_index]
-          when :left_semi, :left_anti
-            [left_index]
-          when :right_semi, :right_anti
+          case type_nick
+          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)
-        dropper = []
-        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
-        dataframe = dataframe.drop(dropper)
-        sorter = [left_index, right_index]
-      when :right_outer
-        dataframe =
-          if joined_table.keys.uniq!
-            DataFrame.create(rename_table(joined_table, left_outputs.size, suffix))
+          when 'left-semi', 'left-anti'
+            [left_index]
           else
-            DataFrame.create(joined_table)
+            [left_index, right_index]
           end
-        dataframe = dataframe.pick(right_keys, dataframe.keys - right_keys)
-        sorter = [left_index, right_index]
-      end
-
-      if force_order
-        dataframe
-          .sort(sorter)
+        df.sort(sorter)
           .drop(sorter)
       else
-        dataframe
+        df
       end
     end
 
     private
 
-    # To ensure Array of Symbols
+    # To ensure Array of Strings
     def ensure_keys(keys)
-      Array(keys).map(&:to_sym)
+      Array(keys).map(&:to_s)
+    end
+
+    # Merge key columns and preserve as left and remove right.
+    def merge_keys(plan, input_node, context)
+      type_nick, left_table_keys, right_table_keys, left_keys, right_keys, * = context
+      return input_node unless type_nick == 'full-outer'
+
+      left_indices = left_keys.map { left_table_keys.index(_1) }
+      right_offset = left_table_keys.size
+      right_indices = right_keys.map { right_table_keys.index(_1) + right_offset }
+      expressions = []
+      names = []
+      left_table_keys.each_with_index do |key, index|
+        names << key
+        expressions <<
+          if (i = left_indices.index(index))
+            left_field = Arrow::FieldExpression.new("[#{left_indices[i]}]")
+            right_field = Arrow::FieldExpression.new("[#{right_indices[i]}]")
+            is_left_null = Arrow::CallExpression.new('is_null', [left_field])
+            Arrow::CallExpression.new('if_else', [is_left_null, right_field, left_field])
+          else
+            Arrow::FieldExpression.new("[#{index}]")
+          end
+      end
+      right_table_keys.each.with_index(right_offset) do |key, index|
+        unless right_indices.include?(index)
+          names << key
+          expressions << Arrow::FieldExpression.new("[#{index}]")
+        end
+      end
+      project_node_options = Arrow::ProjectNodeOptions.new(expressions, names)
+      plan.build_project_node(input_node, project_node_options)
     end
 
-    # Rename duplicate keys by suffix
-    def rename_table(joined_table, n_keys, suffix)
-      joined_keys = joined_table.keys
-      other_keys = joined_keys[n_keys..]
+    def rename_keys(plan, input_node, context)
+      type_nick, left_table_keys, right_table_keys, *, suffix = context
+      names = input_node.output_schema.fields.map(&:name)
+      return input_node unless names.dup.uniq!
 
-      dup_keys = joined_keys.tally.select { |_, v| v > 1 }.keys
+      pos_rights =
+        if type_nick.start_with?('right')
+          names.size - right_table_keys.size
+        else
+          left_table_keys.size
+        end
+      rights = names[pos_rights..]
+      dup_keys = names.tally.select { |_, v| v > 1 }.keys
       renamed_right_keys =
-        other_keys.map do |key|
+        rights.map do |key|
           if dup_keys.include?(key)
-            suffixed = "#{key}#{suffix}".to_sym
+            suffixed = "#{key}#{suffix}".to_s
             # Find a key from suffixed.succ
-            (suffixed..).find { !joined_keys.include?(_1) }
+            (suffixed..).find { !names.include?(_1) }
           else
             key
           end
         end
-      joined_keys[n_keys..] = renamed_right_keys
-
-      fields =
-        joined_keys.map.with_index do |k, i|
-          Arrow::Field.new(k, joined_table[i].data_type)
-        end
-      Arrow::Table.new(Arrow::Schema.new(fields), joined_table.columns)
-    end
+      names[pos_rights..] = renamed_right_keys
 
-    # Merge two Arrow::Arrays
-    def merge_array(array1, array2)
-      t = Arrow::Function.find(:is_null).execute([array1])
-      Arrow::Function.find(:if_else).execute([t, array2, array1]).value
+      expressions = names.map.with_index { |_, i| Arrow::FieldExpression.new("[#{i}]") }
+      project_node_options = Arrow::ProjectNodeOptions.new(expressions, names)
+      plan.build_project_node(input_node, project_node_options)
     end
   end
 end

data/lib/red_amber/data_frame_displayable.rb

@@ -269,12 +269,13 @@ module RedAmber
     end
     alias_method :glimpse, :tdr
 
-    # Shortcut for `tdr(:all)``.
+    # Shortcut for `tdr(:all)`.
     #
+    # @param (see #tdr)
     # @return (see #tdr)
     #
-    def tdra
-      puts tdr_str(:all)
+    def tdra(tally: 5, elements: 5)
+      puts tdr_str(:all, tally: tally, elements: elements)
     end
 
     # rubocop:enable Layout/LineLength
@@ -504,9 +505,9 @@ module RedAmber
           row.zip(formats).map do |elem, format|
             non_ascii_diff = elem.ascii_only? ? 0 : elem.width - elem.size
             if format.negative?
-              elem.ljust(-format + non_ascii_diff)
+              elem.ljust(-format - non_ascii_diff)
             else
-              elem.rjust(format + non_ascii_diff)
+              elem.rjust(format - non_ascii_diff)
             end
           end
         str.puts a.join(' ').rstrip

data/lib/red_amber/data_frame_selectable.rb

@@ -836,6 +836,55 @@ module RedAmber
       tail(n_obs)
     end
 
+    # Select records randomly to create a DataFrame.
+    #   This method calls `indices.sample`.
+    #   We can use the same arguments in `Vector#sample`.
+    # @note This method requires 'arrow-numo-narray' gem.
+    #
+    # @overload sample()
+    #   Return a DataFrame with a randomly selected record.
+    #
+    #   @return [DataFrame]
+    #     a DataFrame with single record.
+    #
+    # @overload sample(n)
+    #   Return a DataFrame with n records selected at random.
+    #
+    #   @param n [Integer]
+    #     positive number of records to select.
+    #     If n is smaller or equal to size, records are selected by non-repeating.
+    #     If n is greater than `size`, records are selected repeatedly.
+    #   @return [DataFrame]
+    #     a DataFrame with sampled records.
+    #
+    # @overload sample(prop)
+    #   Return a DataFrame with records by proportion `prop` at random.
+    #
+    #   @param prop [Float]
+    #     positive proportion of records to select.
+    #     Absolute number of records to select:`prop*size` is rounded (by `half: :up`).
+    #     If prop is smaller or equal to 1.0, records are selected by non-repeating.
+    #     If prop is greater than 1.0, some records are selected repeatedly.
+    #   @return [Vector]
+    #     a DataFrame with sampled records.
+    #
+    # @since 0.5.0
+    #
+    def sample(n_or_prop = nil)
+      slice { indices.sample(n_or_prop) }
+    end
+
+    # Returns a DataFrame with shuffled rows.
+    #
+    # @note This method requires 'arrow-numo-narray' gem.
+    # @note Same behavior as `DataFrame#sample(1.0)`
+    # @return (see #sample)
+    # @since 0.5.0
+    #
+    def shuffle
+      sample(1.0)
+    end
+
     # Select records by index Array to create a DataFrame.
     #
     # - TODO: support for option `boundscheck: true`