red_amber 0.2.2 → 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?
-          picker
-        else
-          raise DataFrameArgumentError, "Invalid argument #{args}"
-        end
 
-      # DataFrame#[] creates a Vector with 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 - dropper
+      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 with 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

@@ -3,35 +3,88 @@
 module RedAmber
   # group class
   class Group
+    include Enumerable # This feature is experimental
+
+    using RefineArrowTable
+
     # Creates a new Group object.
     #
     # @param dataframe [DataFrame] dataframe to be grouped.
     # @param group_keys [Array<>] keys for grouping.
     def initialize(dataframe, *group_keys)
       @dataframe = dataframe
-      @table = @dataframe.table
       @group_keys = group_keys.flatten
 
-      raise GroupArgumentError, 'group_keys is empty.' if @group_keys.empty?
+      raise GroupArgumentError, 'group_keys are empty.' if @group_keys.empty?
 
       d = @group_keys - @dataframe.keys
       raise GroupArgumentError, "#{d} is not a key of\n #{@dataframe}." unless d.empty?
 
-      @group = @table.group(*@group_keys)
+      @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|
-        by(function, 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
 
-    def inspect
-      tallys = @dataframe.pick(@group_keys).vectors.map.with_object({}) do |v, h|
-        h[v.key] = v.tally
+    alias_method :__count, :count
+    private :__count
+
+    def count(*summary_keys)
+      df = __count(summary_keys)
+      # if counts are the same (and do not include NaN or nil), aggregate count columns.
+      if df.pick(@group_keys.size..).to_h.values.uniq.size == 1
+        df.pick(0..@group_keys.size).rename { [keys[-1], :count] }
+      else
+        df
       end
-      "#<#{self.class}:#{format('0x%016x', object_id)}\n#{tallys}>"
+    end
+
+    def filters
+      @filters ||= begin
+        first, *others = @group_keys.map do |key|
+          vector = @dataframe[key]
+          vector.uniq.each.map { |u| u.nil? ? vector.is_nil : vector == u }
+        end
+
+        if others.empty?
+          first.select(&:any?)
+        else
+          first.product(*others).map { |a| a.reduce(&:&) }.select(&:any?)
+        end
+      end
+    end
+
+    def each
+      filters
+      return enum_for(:each) unless block_given?
+
+      @filters.each do |filter|
+        yield @dataframe[filter]
+      end
+      @filters.size
+    end
+
+    def group_count
+      DataFrame.create(add_columns_to_table(base_table, [:group_count], [group_counts]))
+    end
+
+    def inspect
+      "#<#{self.class} : #{format('0x%016x', object_id)}>\n#{group_count}"
     end
 
     def summarize(&block)
@@ -46,20 +99,57 @@ module RedAmber
       end
     end
 
+    # experimental
+    def agg_sum(*summary_keys)
+      call_aggregating_function(:sum, summary_keys, _options = nil)
+    end
+
     private
 
-    def by(func, 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?
+    def build_aggregation_keys(function_name, summary_keys)
+      if summary_keys.empty?
+        [function_name]
+      else
+        summary_keys.map { |key| "#{function_name}(#{key})" }
+      end
+    end
 
-      df = RedAmber::DataFrame.new(@group.send(func, *summary_keys))
-      df = df.pick(@group_keys, df.keys - @group_keys)
-      # if counts are the same (and do not include NaN or nil), aggregate count columns.
-      if func == :count && df.pick(@group_keys.size..).to_h.values.uniq.size == 1
-        df = df.pick(0..@group_keys.size).rename { [keys[-1], :count] }
+    # @note `@group_counts.sum == @dataframe.size``
+    def group_counts
+      @group_counts ||= filters.map(&:sum)
+    end
+
+    def base_table
+      @base_table ||= begin
+        indexes = filters.map { |filter| filter.index(true) }
+        @dataframe.table[@group_keys].take(indexes)
+      end
+    end
+
+    def add_columns_to_table(table, keys, data_arrays)
+      fields = table.schema.fields
+      arrays = table.columns.map(&:data)
+
+      keys.zip(data_arrays).each do |key, array|
+        data = Arrow::ChunkedArray.new([array])
+        fields << Arrow::Field.new(key, data.value_data_type)
+        arrays << data
+      end
+
+      Arrow::Table.new(Arrow::Schema.new(fields), arrays)
+    end
+
+    # Call Vector aggregating function and return an array of arrays:
+    #   [keys, data_arrays]
+    #   (Experimental feature)
+    def call_aggregating_function(func, summary_keys, _options)
+      summary_keys.each.with_object([[], []]) do |key, (keys, arrays)|
+        vector = @dataframe[key]
+        arrays << filters.map { |filter| vector.filter(filter).send(func) }
+        keys << "#{func}(#{key})".to_sym
+      rescue Arrow::Error::NotImplemented
+        # next
       end
-      df
     end
   end
 end

data/lib/red_amber/helper.rb

@@ -5,58 +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 out_of_range?(indeces)
-      indeces.max >= size || indeces.min < -size
-    end
-
-    def integers?(enum)
-      enum.all?(Integer)
-    end
-
-    def booleans?(enum)
-      enum.all? { |e| e.is_a?(TrueClass) || e.is_a?(FalseClass) || e.is_a?(NilClass) }
-    end
-
-    def create_dataframe_from_vector(key, vector)
-      DataFrame.new(key => vector.data)
-    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