red_amber 0.2.3 → 0.3.0

data/lib/red_amber/data_frame_variable_operation.rb

@@ -3,71 +3,145 @@
 module RedAmber
   # mix-ins for the class DataFrame
   module DataFrameVariableOperation
-    # pick up some variables to create sub DataFrame
+    # Array is refined
+    using RefineArray
+
+    # Pick up variables (columns) to create a new DataFrame
+    #
+    # @note DataFrame#pick creates a DataFrame with single key.
+    #   DataFrame#[] creates a Vector if single key is specified.
+    #
+    # @overload pick(keys)
+    #   Pick variables by Symbols or Strings.
+    #
+    #   @param keys [Symbol, String, <Symbol, String>]
+    #     key name(s) of variables to pick.
+    #   @return [DataFrame]
+    #     Picked DataFrame.
+    #
+    # @overload pick(booleans)
+    #   Pick variables by booleans.
+    #
+    #   @param booleans [<true, false, nil>]
+    #     boolean array to pick variables at true.
+    #   @return [DataFrame]
+    #     Picked DataFrame.
+    #
+    # @overload pick(indices)
+    #   Pick variables by column indices.
+    #
+    #   @param indices [Integer, Float, Range<Integer>, Vector, Arrow::Array]
+    #     numeric array to pick variables by column index.
+    #   @return [DataFrame]
+    #     Picked DataFrame.
+    #
     def pick(*args, &block)
-      picker = args
       if block
-        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+        unless args.empty?
+          raise DataFrameArgumentError, 'Must not specify both arguments and block.'
+        end
 
-        picker = [instance_eval(&block)]
+        args = [instance_eval(&block)]
       end
-      picker.flatten!
-      return DataFrame.new if picker.empty? || picker == [nil]
-
-      key_vector = Vector.new(keys)
-      vec = parse_to_vector(picker, vsize: n_keys)
-
-      ary =
-        if vec.boolean?
-          key_vector.filter(*vec).to_a
-        elsif vec.numeric?
-          key_vector.take(*vec).to_a
-        elsif vec.string? || vec.dictionary?
-          vec.to_a
-        else
-          raise DataFrameArgumentError, "Invalid argument #{args}"
-        end
 
-      # DataFrame#[] creates a Vector if single key is specified.
-      # DataFrame#pick creates a DataFrame with single key.
-      DataFrame.new(@table[ary])
+      case args
+      in [] | [nil]
+        return DataFrame.new
+      in [*] if args.symbols?
+        return DataFrame.create(@table.select_columns(*args))
+      in [*] if args.booleans?
+        picker = keys.select_by_booleans(args)
+        return DataFrame.create(@table.select_columns(*picker))
+      in [(Vector | Arrow::Array | Arrow::ChunkedArray) => a]
+        picker = a.to_a
+      else
+        picker = parse_args(args, n_keys)
+      end
+
+      return DataFrame.new if picker.compact.empty?
+
+      if picker.booleans?
+        picker = keys.select_by_booleans(picker)
+        return DataFrame.create(@table.select_columns(*picker))
+      end
+      picker.compact!
+      raise DataFrameArgumentError, "some keys are duplicated: #{args}" if picker.uniq!
+
+      DataFrame.create(@table.select_columns(*picker))
     end
 
-    # drop some variables to create remainer sub DataFrame
+    # Drop some variables (columns) to create a remainer DataFrame
+    #
+    # @note DataFrame#drop creates a DataFrame even if it is a single column.
+    #
+    # @overload drop(keys)
+    #   Drop variables by Symbols or Strings.
+    #
+    #   @param keys [Symbol, String, <Symbol, String>]
+    #     key name(s) of variables to drop.
+    #   @return [DataFrame]
+    #     Remainer DataFrame.
+    #
+    # @overload drop(booleans)
+    #   Drop variables by booleans.
+    #
+    #   @param booleans [<true, false, nil>]
+    #     boolean array of variables to drop at true.
+    #   @return [DataFrame]
+    #     Remainer DataFrame.
+    #
+    # @overload drop(indices)
+    #   Pick variables by column indices.
+    #
+    #   @param indices [Integer, Float, Range<Integer>, Vector, Arrow::Array]
+    #     numeric array of variables to drop by column index.
+    #   @return [DataFrame]
+    #     Remainer DataFrame.
+    #
     def drop(*args, &block)
-      dropper = args
       if block
-        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+        unless args.empty?
+          raise DataFrameArgumentError, 'Must not specify both arguments and block.'
+        end
 
-        dropper = [instance_eval(&block)]
+        args = [instance_eval(&block)]
       end
-      dropper.flatten!
-
-      key_vector = Vector.new(keys)
-      vec = parse_to_vector(dropper, vsize: n_keys)
-
-      ary =
-        if vec.boolean?
-          key_vector.filter(*vec.primitive_invert).each.map(&:to_sym) # Array
-        elsif vec.numeric?
-          keys - key_vector.take(*vec).each.map(&:to_sym) # Array
-        elsif vec.string? || vec.dictionary?
-          keys - vec.to_a.map { _1&.to_sym } # Array
+      return self if args.empty? || empty?
+
+      picker =
+        if args.symbols?
+          keys - args
+        elsif args.booleans?
+          keys.reject_by_booleans(args)
+        elsif args.integers?
+          keys.reject_by_indices(args)
         else
-          raise DataFrameArgumentError, "Invalid argument #{args}"
+          dropper = parse_args(args, n_keys)
+          if dropper.booleans?
+            keys.reject_by_booleans(dropper)
+          elsif dropper.symbols?
+            keys - dropper
+          else
+            dropper.compact!
+            unless dropper.integers?
+              raise DataFrameArgumentError, "Invalid argument #{args}"
+            end
+
+            keys.reject_by_indices(dropper)
+          end
         end
 
-      return DataFrame.new if ary.empty?
+      return DataFrame.new if picker.empty?
 
-      # DataFrame#[] creates a Vector if single key is specified.
-      # DataFrame#drop creates a DataFrame with single key.
-      DataFrame.new(@table[ary])
+      DataFrame.create(@table.select_columns(*picker))
     end
 
     # rename variables to create a new DataFrame
     def rename(*renamer, &block)
       if block
-        raise DataFrameArgumentError, 'Must not specify both arguments and a block' unless renamer.empty?
+        unless renamer.empty?
+          raise DataFrameArgumentError, 'Must not specify both arguments and a block'
+        end
 
         renamer = [instance_eval(&block)]
       end
@@ -90,35 +164,23 @@ module RedAmber
 
     # assign variables to create a new DataFrame
     def assign(*assigner, &block)
-      appender, fields, arrays = assign_update(*assigner, &block)
-      return self if appender.is_a?(DataFrame)
-
-      append_to_fields_and_arrays(appender, fields, arrays, append_to_left: false) unless appender.empty?
-
-      DataFrame.new(Arrow::Table.new(Arrow::Schema.new(fields), arrays))
+      assign_update(*assigner, append_to_left: false, &block)
     end
 
     def assign_left(*assigner, &block)
-      appender, fields, arrays = assign_update(*assigner, &block)
-      return self if appender.is_a?(DataFrame)
-
-      append_to_fields_and_arrays(appender, fields, arrays, append_to_left: true) unless appender.empty?
-
-      DataFrame.new(Arrow::Table.new(Arrow::Schema.new(fields), arrays))
+      assign_update(*assigner, append_to_left: true, &block)
     end
 
     private
 
-    def assign_update(*assigner, &block)
+    def assign_update(*assigner, append_to_left: false, &block)
       if block
         assigner_from_block = instance_eval(&block)
         assigner =
-          if assigner.empty?
-            # block only
+          case assigner_from_block
+          in _ if assigner.empty? # block only
             [assigner_from_block]
-          # If Ruby >= 3.0, one line pattern match can be used
-          # assigner_from_block in [Array, *]
-          elsif multiple_assigner?(assigner_from_block)
+          in [Vector, *] | [Array, *] | [Arrow::Array, *]
             assigner.zip(assigner_from_block)
           else
             assigner.zip([assigner_from_block])
@@ -128,10 +190,10 @@ module RedAmber
       case assigner
       in [] | [nil] | [{}] | [[]]
         return self
-      in [Hash => key_array_pairs]
-      # noop
       in [(Symbol | String) => key, (Vector | Array | Arrow::Array) => array]
         key_array_pairs = { key => array }
+      in [Hash => key_array_pairs]
+      # noop
       in [Array => array_in_array]
         key_array_pairs = try_convert_to_hash(array_in_array)
       in [Array, *] => array_in_array1
@@ -151,20 +213,27 @@ module RedAmber
           appender[key] = array
         end
       end
-      [appender, *update_fields_and_arrays(updater)]
+      fields, arrays = *update_fields_and_arrays(updater)
+      return self if appender.is_a?(DataFrame)
+
+      unless appender.empty?
+        append_to_fields_and_arrays(appender, fields, arrays, append_to_left)
+      end
+
+      DataFrame.create(Arrow::Table.new(Arrow::Schema.new(fields), arrays))
     end
 
     def try_convert_to_hash(array)
       array.to_h
     rescue TypeError
       [array].to_h
-    rescue TypeError # rubocop:disable Lint/DuplicateRescueException
-      raise DataFrameArgumentError, "Invalid argument in Array #{array}"
     end
 
     def rename_by_hash(key_pairs)
       not_existing_keys = key_pairs.keys - keys
-      raise DataFrameArgumentError, "Not existing: #{not_existing_keys}" unless not_existing_keys.empty?
+      unless not_existing_keys.empty?
+        raise DataFrameArgumentError, "Not existing: #{not_existing_keys}"
+      end
 
       fields =
         keys.map do |key|
@@ -175,7 +244,7 @@ module RedAmber
             @table.schema[key]
           end
         end
-      DataFrame.new(Arrow::Table.new(Arrow::Schema.new(fields), @table.columns))
+      DataFrame.create(Arrow::Table.new(Arrow::Schema.new(fields), @table.columns))
     end
 
     def update_fields_and_arrays(updater)
@@ -185,7 +254,9 @@ module RedAmber
         data = updater[key]
         next unless data
 
-        raise DataFrameArgumentError, "Data size mismatch (#{data.size} != #{size})" if data.nil? || data.size != size
+        if data.size != size
+          raise DataFrameArgumentError, "Data size mismatch (#{data.size} != #{size})"
+        end
 
         a = Arrow::Array.new(data.is_a?(Vector) ? data.to_a : data)
         fields[i] = Arrow::Field.new(key, a.value_data_type)
@@ -194,10 +265,12 @@ module RedAmber
       [fields, arrays]
     end
 
-    def append_to_fields_and_arrays(appender, fields, arrays, append_to_left: false)
+    def append_to_fields_and_arrays(appender, fields, arrays, append_to_left)
       enum = append_to_left ? appender.reverse_each : appender.each
       enum.each do |key, data|
-        raise DataFrameArgumentError, "Data size mismatch (#{data.size} != #{size})" if data.size != size
+        if data.size != size
+          raise DataFrameArgumentError, "Data size mismatch (#{data.size} != #{size})"
+        end
 
         a = Arrow::Array.new(data.is_a?(Vector) ? data.to_a : data)
 
@@ -210,14 +283,5 @@ module RedAmber
         end
       end
     end
-
-    def multiple_assigner?(assigner)
-      case assigner
-      in [Vector, *] | [Array, *] | [Arrow::Array, *]
-        true
-      else
-        false
-      end
-    end
   end
 end

data/lib/red_amber/group.rb

@@ -5,6 +5,8 @@ module RedAmber
   class Group
     include Enumerable # This feature is experimental
 
+    using RefineArrowTable
+
     # Creates a new Group object.
     #
     # @param dataframe [DataFrame] dataframe to be grouped.
@@ -18,7 +20,6 @@ module RedAmber
       d = @group_keys - @dataframe.keys
       raise GroupArgumentError, "#{d} is not a key of\n #{@dataframe}." unless d.empty?
 
-      @filters = @group_counts = @base_table = nil
       @group = @dataframe.table.group(*@group_keys)
     end
 
@@ -29,11 +30,14 @@ module RedAmber
       define_method(function) do |*summary_keys|
         summary_keys = Array(summary_keys).flatten
         d = summary_keys - @dataframe.keys
-        raise GroupArgumentError, "#{d} is not a key of\n #{@dataframe}." unless summary_keys.empty? || d.empty?
+        unless summary_keys.empty? || d.empty?
+          raise GroupArgumentError, "#{d} is not a key of\n #{@dataframe}."
+        end
 
-        table = @group.aggregate(*build_aggregation_keys("hash_#{function}", summary_keys))
-        df = DataFrame.new(table)
-        df.pick(@group_keys, df.keys - @group_keys)
+        table = @group.aggregate(*build_aggregation_keys("hash_#{function}",
+                                                         summary_keys))
+        g = @group_keys.map(&:to_s)
+        DataFrame.new(table[g + (table.keys - g)])
       end
     end
 
@@ -76,7 +80,7 @@ module RedAmber
     end
 
     def group_count
-      DataFrame.new(add_columns_to_table(base_table, [:group_count], [group_counts]))
+      DataFrame.create(add_columns_to_table(base_table, [:group_count], [group_counts]))
     end
 
     def inspect
@@ -95,6 +99,11 @@ module RedAmber
       end
     end
 
+    # experimental
+    def agg_sum(*summary_keys)
+      call_aggregating_function(:sum, summary_keys, _options = nil)
+    end
+
     private
 
     def build_aggregation_keys(function_name, summary_keys)
@@ -105,7 +114,7 @@ module RedAmber
       end
     end
 
-    # @group_counts.sum == @dataframe.size
+    # @note `@group_counts.sum == @dataframe.size``
     def group_counts
       @group_counts ||= filters.map(&:sum)
     end

data/lib/red_amber/helper.rb

@@ -5,46 +5,68 @@ module RedAmber
   module Helper
     private
 
+    # If num is larger than 1 return 's' to be plural.
+    #
+    # @param num [Numeric] some number.
+    # @return ['s', ''] return 's' if num is larger than 1.
+    #   Otherwise return ''.
     def pl(num)
       num > 1 ? 's' : ''
     end
 
-    def booleans?(enum)
-      enum.all? { |e| e.is_a?(TrueClass) || e.is_a?(FalseClass) || e.is_a?(NilClass) }
-    end
-
-    def parse_to_vector(args, vsize: size)
-      a = args.reduce([]) do |accum, elem|
-        accum.concat(normalize_element(elem, vsize: vsize))
+    # Parse the argments in an Array
+    #   and returns a parsed Array.
+    #
+    # @param args
+    #   [<Integer, Symbol, true, false, nil, Array, Range, Enumerator, String, Float>]
+    #   arguments.
+    # @param array_size [Integer] size of target Array to use in a endless Range.
+    # @return [<Integer, Symbol, true, false, nil>] parsed flat Array.
+    # @note This method is recursively called to parse.
+    def parse_args(args, array_size)
+      args.flat_map do |elem|
+        case elem
+        when Integer, Symbol, NilClass, TrueClass, FalseClass
+          elem
+        when Array
+          parse_args(elem, array_size)
+        when Range
+          parse_range(elem, array_size)
+        when Enumerator
+          parse_args(Array(elem), array_size)
+        when String
+          elem.to_sym
+        when Float
+          elem.floor.to_i
+        else
+          Array(elem)
+        end
       end
-      Vector.new(a)
     end
 
-    def normalize_element(elem, vsize: size)
-      case elem
-      when NilClass
-        [nil]
-      when Range
-        bg = elem.begin
-        en = elem.end
-        if [bg, en].any?(Integer)
-          bg += vsize if bg&.negative?
-          en += vsize if en&.negative?
-          en -= 1 if en.is_a?(Integer) && elem.exclude_end?
-          if bg&.negative? || (en && en >= vsize)
-            raise DataFrameArgumentError, "Index out of range: #{elem} for 0..#{vsize - 1}"
-          end
-
-          Array(0...vsize)[elem]
-        elsif bg.nil? && en.nil?
-          Array(0...vsize)
-        else
-          Array(elem)
+    # Parse a Range to an Array
+    #
+    # @param range [Range] Range to parse.
+    # @param array_size [Integer] size of target Array to use in a endless Range.
+    # @return [Array<Integer, Symbol, String>] parsed Array.
+    def parse_range(range, array_size)
+      bg = range.begin
+      en = range.end
+      if [bg, en].any?(Integer)
+        bg += array_size if bg&.negative?
+        en += array_size if en&.negative?
+        en -= 1 if en.is_a?(Integer) && range.exclude_end?
+        if bg&.negative? || (en && en >= array_size)
+          raise IndexError, "Index out of range: #{range} for 0..#{array_size - 1}"
         end
-      when Enumerator
-        elem.to_a
+
+        Array(0...array_size)[range]
+      elsif bg.nil?
+        raise DataFrameArgumentError, "Cannot use beginless Range: #{range}"
+      elsif en.nil?
+        raise DataFrameArgumentError, "Cannot use endless Range: #{range}"
       else
-        Array[elem]
+        Array(range)
       end
     end
   end

data/lib/red_amber/refinements.rb

@@ -0,0 +1,199 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # Add additional capabilities to Hash
+  module RefineHash
+    refine Hash do
+      # Convert self to an Arrow::Table
+      def to_arrow
+        Arrow::Table.new(self)
+      end
+    end
+  end
+
+  # Add additional capabilities to Array-like classes
+  module RefineArrayLike
+    refine Array do
+      def to_arrow_array
+        Arrow::Array.new(self)
+      end
+    end
+
+    refine Range do
+      def to_arrow_array
+        Arrow::Array.new(Array(self))
+      end
+    end
+
+    # common methods for Arrow::Array and Arrow::ChunkedArray
+    # Refinement#include is deprecated and will be removed in Ruby 3.2
+    refine Arrow::Array do
+      def to_arrow_array
+        self
+      end
+
+      def type_class
+        value_data_type.class
+      end
+
+      def boolean?
+        value_data_type.instance_of?(Arrow::BooleanDataType)
+      end
+
+      def numeric?
+        value_data_type.class < Arrow::NumericDataType
+      end
+
+      def float?
+        value_data_type.class < Arrow::FloatingPointDataType
+      end
+
+      def integer?
+        value_data_type.class < Arrow::IntegerDataType
+      end
+
+      def list?
+        is_a? Arrow::ListArray
+      end
+
+      def unsigned_integer?
+        value_data_type.instance_of?(Arrow::UInt8DataType) ||
+          value_data_type.instance_of?(Arrow::UInt16DataType) ||
+          value_data_type.instance_of?(Arrow::UInt32DataType) ||
+          value_data_type.instance_of?(Arrow::UInt64DataType)
+      end
+
+      def string?
+        value_data_type.instance_of?(Arrow::StringDataType)
+      end
+
+      def dictionary?
+        value_data_type.instance_of?(Arrow::DictionaryDataType)
+      end
+
+      def temporal?
+        value_data_type.class < Arrow::TemporalDataType
+      end
+
+      def primitive_invert
+        n = Arrow::Function.find(:is_null).execute([self])
+        i = Arrow::Function.find(:if_else).execute([n, false, self])
+        Arrow::Function.find(:invert).execute([i]).value
+      end
+    end
+
+    refine Arrow::ChunkedArray do
+      def to_arrow_array
+        self
+      end
+
+      def type_class
+        value_data_type.class
+      end
+
+      def boolean?
+        value_data_type.instance_of?(Arrow::BooleanDataType)
+      end
+
+      def numeric?
+        value_data_type.class < Arrow::NumericDataType
+      end
+
+      def float?
+        value_data_type.class < Arrow::FloatingPointDataType
+      end
+
+      def integer?
+        value_data_type.class < Arrow::IntegerDataType
+      end
+
+      def unsigned_integer?
+        value_data_type.instance_of?(Arrow::UInt8DataType) ||
+          value_data_type.instance_of?(Arrow::UInt16DataType) ||
+          value_data_type.instance_of?(Arrow::UInt32DataType) ||
+          value_data_type.instance_of?(Arrow::UInt64DataType)
+      end
+
+      def string?
+        value_data_type.instance_of?(Arrow::StringDataType)
+      end
+
+      def dictionary?
+        value_data_type.instance_of?(Arrow::DictionaryDataType)
+      end
+
+      def temporal?
+        value_data_type.class < Arrow::TemporalDataType
+      end
+
+      def list?
+        value_type.nick == 'list'
+      end
+
+      def primitive_invert
+        n = Arrow::Function.find(:is_null).execute([self])
+        i = Arrow::Function.find(:if_else).execute([n, false, self])
+        Arrow::Function.find(:invert).execute([i]).value
+      end
+    end
+  end
+
+  # Add additional capabilities to Arrow::Table
+  module RefineArrowTable
+    refine Arrow::Table do
+      def keys
+        columns.map(&:name)
+      end
+
+      def key?(key)
+        keys.include?(key)
+      end
+    end
+  end
+
+  # Add additional capabilities to Array
+  module RefineArray
+    refine Array do
+      def integers?
+        all? { |e| e.is_a?(Integer) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def booleans?
+        all? { |e| e.is_a?(TrueClass) || e.is_a?(FalseClass) || e.is_a?(NilClass) }
+      end
+
+      def symbols?
+        all? { |e| e.is_a?(Symbol) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def strings?
+        all? { |e| e.is_a?(String) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def symbols_or_strings?
+        all? { |e| e.is_a?(Symbol) || e.is_a?(String) }
+      end
+
+      # convert booleans to indices
+      def booleans_to_indices
+        (0...size).select.with_index { |_, i| self[i] }
+      end
+
+      # select elements by booleans
+      def select_by_booleans(booleans)
+        select.with_index { |_, i| booleans[i] }
+      end
+
+      # reject elements by booleans
+      def reject_by_booleans(booleans)
+        reject.with_index { |_, i| booleans[i] }
+      end
+
+      # reject elements by indices
+      # notice: order by indices is not considered.
+      def reject_by_indices(indices)
+        reject.with_index { |_, i| indices.include?(i) || indices.include?(i - size) }
+      end
+    end
+  end
+end