red_amber 0.2.3 → 0.4.0

data/lib/red_amber/group.rb

@@ -1,14 +1,72 @@
 # frozen_string_literal: true
 
 module RedAmber
-  # group class
+  # Group class
   class Group
     include Enumerable # This feature is experimental
 
+    using RefineArrowTable
+
+    # Source DataFrame.
+    #
+    # @return [DataFrame]
+    #   source DataFrame.
+    #
+    attr_reader :dataframe
+
+    # Keys for grouping by value.
+    #
+    # @return [Array]
+    #   group keys.
+    #
+    attr_reader :group_keys
+
+    class << self
+      private
+
+      # @!macro [attach] define_group_aggregation
+      #   @!method $1(*summary_keys)
+      #     Group aggregation function `$1`.
+      #     @param summary_keys [Array<Symbol, String>]
+      #       summary keys.
+      #     @return [DataFrame]
+      #       aggregated DataFrame
+      #
+      def define_group_aggregation(function)
+        define_method(function) do |*summary_keys|
+          summary_keys = Array(summary_keys).flatten
+          d = summary_keys - @dataframe.keys
+          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))
+          g = @group_keys.map(&:to_s)
+          DataFrame.new(table[g + (table.keys - g)])
+        end
+      end
+    end
+
     # Creates a new Group object.
     #
-    # @param dataframe [DataFrame] dataframe to be grouped.
-    # @param group_keys [Array<>] keys for grouping.
+    # @param dataframe [DataFrame]
+    #   dataframe to be grouped.
+    # @param group_keys [Array<Symbol, String>]
+    #   keys for grouping.
+    # @return [Group]
+    #   Group object.
+    # @example
+    #   Group.new(penguins, :species)
+    #
+    #   # =>
+    #   #<RedAmber::Group : 0x000000000000f410>
+    #     species   group_count
+    #     <string>      <uint8>
+    #   0 Adelie            152
+    #   1 Chinstrap          68
+    #   2 Gentoo            124
+    #
     def initialize(dataframe, *group_keys)
       @dataframe = dataframe
       @group_keys = group_keys.flatten
@@ -18,25 +76,10 @@ 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
 
-    attr_reader :dataframe, :group_keys
-
-    functions = %i[count sum product mean min max stddev variance]
-    functions.each do |function|
-      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?
-
-        table = @group.aggregate(*build_aggregation_keys("hash_#{function}", summary_keys))
-        df = DataFrame.new(table)
-        df.pick(@group_keys, df.keys - @group_keys)
-      end
-    end
-
+    define_group_aggregation(:count)
     alias_method :__count, :count
     private :__count
 
@@ -50,6 +93,26 @@ module RedAmber
       end
     end
 
+    define_group_aggregation(:sum)
+
+    define_group_aggregation(:product)
+
+    define_group_aggregation(:mean)
+
+    define_group_aggregation(:min)
+
+    define_group_aggregation(:max)
+
+    define_group_aggregation(:stddev)
+
+    define_group_aggregation(:variance)
+
+    # Returns Array of boolean filters to select each records in the Group.
+    #
+    # @api private
+    # @return [Array]
+    #   an Array of boolean filter Vectors.
+    #
     def filters
       @filters ||= begin
         first, *others = @group_keys.map do |key|
@@ -65,6 +128,25 @@ module RedAmber
       end
     end
 
+    # Iterates over each record group as a DataFrame or returns a Enumerator.
+    #
+    # @api private
+    # @overload each
+    #   Returns a new Enumerator if no block given.
+    #
+    #   @return [Enumerator]
+    #     Enumerator of each group as a DataFrame.
+    #
+    # @overload each
+    #   When a block given, passes each record group as a DataFrame to the block.
+    #
+    #   @yieldparam df [DataFrame]
+    #     passes each record group as a DataFrame by a block parameter.
+    #   @yieldreturn [Object]
+    #     evaluated result value from the block.
+    #   @return [Integer]
+    #     group size.
+    #
     def each
       filters
       return enum_for(:each) unless block_given?
@@ -75,14 +157,98 @@ module RedAmber
       @filters.size
     end
 
+    # Returns each record group size as a DataFrame.
+    #
+    # @return [DataFrame]
+    #   DataFrame consists of:
+    #   - Group key columns.
+    #   - Result columns by group aggregation.
+    # @example
+    #   penguins.group(:species).group_count
+    #
+    #   # =>
+    #   #<RedAmber::DataFrame : 3 x 2 Vectors, 0x0000000000003a70>
+    #     species   group_count
+    #     <string>      <uint8>
+    #   0 Adelie            152
+    #   1 Chinstrap          68
+    #   2 Gentoo            124
+    #
     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
 
+    # String representation of self.
+    #
+    # @return [String]
+    #   show information of self as a String.
+    # @example
+    #   puts penguins.group(:species).inspect
+    #
+    #   # =>
+    #   #<RedAmber::Group : 0x0000000000003a98>
+    #     species   group_count
+    #     <string>      <uint8>
+    #   0 Adelie            152
+    #   1 Chinstrap          68
+    #   2 Gentoo            124
+    #
     def inspect
       "#<#{self.class} : #{format('0x%016x', object_id)}>\n#{group_count}"
     end
 
+    # Summarize Group by aggregation functions from the block.
+    #
+    # @yieldparam group [Group]
+    #   passes group object self.
+    # @yieldreturn [DataFrame, Array<DataFrame>]
+    #   an aggregated DataFrame or an array of aggregated DataFrames.
+    # @return [DataFrame]
+    #   summarized DataFrame.
+    # @example Single function and single variable
+    #   group = penguins.group(:species)
+    #   group
+    #
+    #   # =>
+    #   #<RedAmber::Group : 0x000000000000c314>
+    #     species   group_count
+    #     <string>      <uint8>
+    #   0 Adelie            152
+    #   1 Chinstrap          68
+    #   2 Gentoo            124
+    #
+    #   group.summarize { mean(:bill_length_mm) }
+    #
+    #   # =>
+    #   #<RedAmber::DataFrame : 3 x 2 Vectors, 0x000000000000c364>
+    #     species   mean(bill_length_mm)
+    #     <string>              <double>
+    #   0 Adelie                   38.79
+    #   1 Chinstrap                48.83
+    #   2 Gentoo                    47.5
+    #
+    # @example Single function only
+    #   group.summarize { mean }
+    #
+    #   # =>
+    #   #<RedAmber::DataFrame : 3 x 6 Vectors, 0x000000000000c350>
+    #     species   mean(bill_length_mm) mean(bill_depth_mm) ... mean(year)
+    #     <string>              <double>            <double> ...   <double>
+    #   0 Adelie                   38.79               18.35 ...    2008.01
+    #   1 Chinstrap                48.83               18.42 ...    2007.97
+    #   2 Gentoo                    47.5               14.98 ...    2008.08
+    #
+    # @example Multiple functions
+    #   group.summarize { [min(:bill_length_mm), max(:bill_length_mm)] }
+    #
+    #   # =>
+    #   #<RedAmber::DataFrame : 3 x 3 Vectors, 0x000000000000c378>
+    #     species   min(bill_length_mm) max(bill_length_mm)
+    #     <string>             <double>            <double>
+    #   0 Adelie                   32.1                46.0
+    #   1 Chinstrap                40.9                58.0
+    #   2 Gentoo                   40.9                59.6
+    #
     def summarize(&block)
       agg = instance_eval(&block)
       case agg
@@ -95,6 +261,14 @@ module RedAmber
       end
     end
 
+    # Aggregating summary.
+    #
+    # @api private
+    #
+    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 +279,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

@@ -1,51 +1,133 @@
 # frozen_string_literal: true
 
 module RedAmber
-  # mix-in for the class DataFrame
+  # Mix-in for the class DataFrame
   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
+
+  # rubocop:disable Layout/LineLength
+
+  # Helper for Arrow Functions
+  module ArrowFunction
+    module_function
+
+    # Find Arrow's compute function.
+    #
+    # {https://arrow.apache.org/docs/cpp/compute.html}
+    # @param function_name [Symbol]
+    #   function name.
+    # @return [Arrow::Function]
+    #   arrow compute function object.
+    # @example
+    #   RedAmber::ArrowFunction.find(:array_sort_indices)
+    #
+    #   # =>
+    #   #<Arrow::Function:0x7fa8838a0d80 ptr=0x7fa87e9b7320 array_sort_indices(array, {order=Ascending, null_placement=AtEnd}): Return the indices that would sort an array>
+    #
+    def find(function_name)
+      Arrow::Function.find(function_name)
+    end
+
+    # Show document of Arrow's compute function.
+    #
+    # @param function_name [Symbol]
+    #   function name.
+    # @return [String]
+    #   document of compute function object.
+    # @example
+    #   puts RedAmber::ArrowFunction.arrow_doc(:array_sort_indices)
+    #
+    #   # =>
+    #   array_sort_indices(array, {order=Ascending, null_placement=AtEnd}): Return the indices that would sort an array
+    #   ------------------
+    #   This function computes an array of indices that define a stable sort
+    #   of the input array.  By default, Null values are considered greater
+    #   than any other value and are therefore sorted at the end of the array.
+    #   For floating-point types, NaNs are considered greater than any
+    #   other non-null value, but smaller than null values.
+    #
+    #   The handling of nulls and NaNs can be changed in ArraySortOptions.
+    #
+    def arrow_doc(function_name)
+      f = find(function_name)
+      "#{f}\n#{'-' * function_name.size}\n#{f.doc.description}"
+    end
+  end
+
+  # rubocop:enable Layout/LineLength
 end

data/lib/red_amber/refinements.rb

@@ -0,0 +1,206 @@
+# frozen_string_literal: true
+
+# Namespace of RedAmber
+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 integer?
+        all? { |e| e.is_a?(Integer) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def numeric?
+        all? { |e| e.is_a?(Numeric) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def boolean?
+        all? { |e| e.is_a?(TrueClass) || e.is_a?(FalseClass) || e.is_a?(NilClass) }
+      end
+
+      def symbol?
+        all? { |e| e.is_a?(Symbol) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def string?
+        all? { |e| e.is_a?(String) } # rubocop:disable Performance/RedundantEqualityComparisonBlock
+      end
+
+      def symbol_or_string?
+        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
+
+  private_constant :RefineArray, :RefineArrayLike, :RefineArrowTable, :RefineHash
+end