red_amber 0.2.2 → 0.2.3

data/lib/red_amber/data_frame.rb

@@ -5,6 +5,7 @@ module RedAmber
   # Variable @table holds an Arrow::Table object.
   class DataFrame
     # mix-in
+    include DataFrameCombinable
     include DataFrameDisplayable
     include DataFrameIndexable
     include DataFrameLoadSave
@@ -47,8 +48,6 @@ module RedAmber
         @table = table
       in [Arrow::Table => table]
         @table = table
-      in [DataFrame => dataframe]
-        @table = dataframe.table
       in [rover_or_hash]
         begin
           # Accepts Rover::DataFrame or Hash
@@ -77,8 +76,9 @@ module RedAmber
     def size
       @table.n_rows
     end
-    alias_method :n_rows, :size
+    alias_method :n_records, :size
     alias_method :n_obs, :size
+    alias_method :n_rows, :size
 
     # Returns the number of columns.
     #
@@ -86,8 +86,9 @@ module RedAmber
     def n_keys
       @table.n_columns
     end
-    alias_method :n_cols, :n_keys
+    alias_method :n_variables, :n_keys
     alias_method :n_vars, :n_keys
+    alias_method :n_cols, :n_keys
 
     # Returns the numbers of rows and columns.
     #
@@ -174,7 +175,7 @@ module RedAmber
     #   - indices(1) #=> [1, 2, 3, 4, 5]
     #   - indices('a') #=> ['a', 'b', 'c', 'd', 'e']
     def indices(start = 0)
-      (start..).take(size)
+      Vector.new((start..).take(size))
     end
     alias_method :indexes, :indices
 

data/lib/red_amber/data_frame_combinable.rb

@@ -0,0 +1,283 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-in for the class DataFrame
+  module DataFrameCombinable
+    # Concatenate other dataframe onto the bottom.
+    #
+    # @param other [DataFrame, Arrow::Table, Array<DataFrame, Arrow::Table>]
+    #   DataFrame/Table to concatenate onto the bottom of self.
+    # @return [DataFrame]
+    #   Concatenated dataframe.
+    def concatenate(*other)
+      case other
+      in [] | [nil] | [[]]
+        return self
+      in [Array => array]
+        # Nop
+      else
+        array = other
+      end
+
+      table_array = array.map do |e|
+        case e
+        when Arrow::Table
+          e
+        when DataFrame
+          e.table
+        else
+          raise DataFrameArgumentError, "#{e} is not a Table or a DataFrame"
+        end
+      end
+
+      DataFrame.new(table.concatenate(table_array))
+    end
+
+    alias_method :concat, :concatenate
+    alias_method :bind_rows, :concatenate
+
+    # Merge other DataFrame or Table from other.
+    # - Self and other must have same size.
+    # - Self and other do not share the same key.
+    #   - If they share any keys, raise Error.
+    # @param other [DataFrame, Arrow::Table, Array<DataFrame, Arrow::Table>]
+    #   DataFrame/Table to concatenate.
+    # @return [DataFrame]
+    #   Merged dataframe.
+    def merge(*other)
+      case other
+      in [] | [nil] | [[]]
+        return self
+      in [Array => array]
+        # Nop
+      else
+        array = other
+      end
+
+      hash = array.each_with_object({}) do |e, h|
+        df =
+          case e
+          when Arrow::Table
+            DataFrame.new(e)
+          when DataFrame
+            e
+          else
+            raise DataFrameArgumentError, "#{e} is not a Table or a DataFrame"
+          end
+
+        raise DataFrameArgumentError, "#{e} do not have same size as self" if size != df.size
+
+        k = keys.intersection(df.keys).any?
+        raise DataFrameArgumentError, "There are some shared keys: #{k}" if k
+
+        h.merge!(df.to_h)
+      end
+
+      assign(hash)
+    end
+
+    alias_method :bind_cols, :merge
+
+    # Mutating joins
+
+    # Join data, leaving only the matching records.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def inner_join(other, join_keys = nil, suffix: '.1')
+      join(other, join_keys, type: :inner, suffix: suffix)
+    end
+
+    # Join data, leaving all records.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def full_join(other, join_keys = nil, suffix: '.1')
+      join(other, join_keys, type: :full_outer, suffix: suffix)
+    end
+
+    alias_method :outer_join, :full_join
+
+    # Join matching values to self from other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def left_join(other, join_keys = nil, suffix: '.1')
+      join(other, join_keys, type: :left_outer, suffix: suffix)
+    end
+
+    # Join matching values from self to other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def right_join(other, join_keys = nil, suffix: '.1')
+      join(other, join_keys, type: :right_outer, suffix: suffix)
+    end
+
+    # Filtering joins
+
+    # Return records of self that have a match in other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def semi_join(other, join_keys = nil, suffix: '.1')
+      join(other, join_keys, type: :left_semi, suffix: suffix)
+    end
+
+    # Return records of self that do not have a match in other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def anti_join(other, join_keys = nil, suffix: '.1')
+      join(other, join_keys, type: :left_anti, suffix: suffix)
+    end
+
+    # Set operations
+
+    # Check if set operation with self and other is possible.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be checked with self.
+    # @return [Boolean] true if set operation is possible.
+    #
+    def set_operable?(other) # rubocop:disable Naming/AccessorMethodName
+      other = DataFrame.new(other) if other.is_a?(Arrow::Table)
+      keys == other.keys
+    end
+
+    # Select records appearing in both self and other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def intersect(other)
+      other = DataFrame.new(other) if other.is_a?(Arrow::Table)
+      raise DataFrameArgumentError, 'keys are not same with self and other' unless keys == other.keys
+
+      join(other, keys, type: :inner)
+    end
+
+    # Select records appearing in self or other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def union(other)
+      other = DataFrame.new(other) if other.is_a?(Arrow::Table)
+      raise DataFrameArgumentError, 'keys are not same with self and other' unless keys == other.keys
+
+      join(other, keys, type: :full_outer)
+    end
+
+    # Select records appearing in self but not in other.
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @return [DataFrame] Joined dataframe.
+    #
+    def difference(other)
+      other = DataFrame.new(other) if other.is_a?(Arrow::Table)
+      raise DataFrameArgumentError, 'keys are not same with self and other' unless keys == other.keys
+
+      join(other, keys, type: :left_anti)
+    end
+
+    alias_method :setdiff, :difference
+
+    # Undocumented. It is preferable to call specific methods.
+
+    # Join other dataframe
+    #
+    # @param other [DataFrame, Arrow::Table] DataFrame/Table to be joined with self.
+    # @param join_keys [String, Symbol, ::Array<String, Symbol>] Keys to match.
+    # @return [DataFrame] Joined dataframe.
+    #
+    #   :type is one of
+    #     :left_semi, :right_semi, :left_anti, :right_anti inner, :left_outer, :right_outer, :full_outer.
+    def join(other, join_keys = nil, type: :inner, suffix: '.1', left_outputs: nil, right_outputs: nil)
+      case other
+      when DataFrame
+        # Nop
+      when Arrow::Table
+        other = DataFrame.new(other)
+      else
+        raise DataFrameArgumentError, 'other must be a DataFrame or an Arrow::Table'
+      end
+
+      # Support natural keys (implicit common keys)
+      natural_keys = keys.intersection(other.keys)
+      raise DataFrameArgumentError, "#{join_keys} are not common keys" if natural_keys.empty?
+
+      join_keys =
+        if join_keys
+          Array(join_keys).map(&:to_sym)
+        else
+          natural_keys
+        end
+      return self if join_keys.empty?
+
+      # Support partial join_keys (common key other than join_key will be renamed with suffix)
+      remainer_keys = natural_keys - join_keys
+      unless remainer_keys.empty?
+        renamer = remainer_keys.each_with_object({}) do |key, hash|
+          new_key = nil
+          loop do
+            new_key = "#{key}#{suffix}".to_sym
+            break unless keys.include?(new_key)
+
+            s = suffix.succ
+            raise DataFrameArgumentError, "suffix #{suffix} is invalid" if s == suffix
+
+            suffix = s
+          end
+          hash[key] = new_key
+        end
+        other = other.rename(renamer)
+      end
+
+      # Red Arrow's #join returns duplicated join_keys from self and other as of v9.0.0 .
+      # Temporally merge key vectors here to workaround.
+      table_output =
+        table.join(other.table, join_keys, type: type, left_outputs: left_outputs, right_outputs: right_outputs)
+      left_indexes = [*0...n_keys]
+      right_indexes = [*((other.keys - join_keys).map { |key| other.keys.index(key) + n_keys })]
+
+      case type
+      when :left_semi, :left_anti, :right_semi, :right_anti
+        return DataFrame.new(table_output)
+      else
+        selected_indexes = left_indexes.concat(right_indexes)
+      end
+      merged_columns = join_keys.map do |key|
+        i = keys.index(key)
+        merge_column(table_output[i], table_output[n_keys + i], type)
+      end
+      DataFrame.new(table_output[selected_indexes])
+               .assign(*join_keys) { merged_columns }
+    end
+
+    private
+
+    def merge_column(column1, column2, type)
+      a1 = column1.to_a
+      a2 = column2.to_a
+      if type == :full_outer
+        a1.zip(a2).map { |x, y| x || y }
+      elsif type.start_with?('right')
+        a2
+      else # :inner or :left-*
+        a1
+      end
+    end
+  end
+end

data/lib/red_amber/data_frame_displayable.rb

@@ -174,6 +174,8 @@ module RedAmber
     end
 
     def format_table(width: 80, head: 5, tail: 3, n_digit: 2)
+      return "  #{keys.join(' ')}\n  (Empty Vectors)\n" if size.zero?
+
       original = self
       indices = size > head + tail ? [*0..head, *(size - tail)...size] : [*0...size]
       df = slice(indices).assign do

data/lib/red_amber/data_frame_selectable.rb

@@ -17,7 +17,7 @@ module RedAmber
         raise DataFrameArgumentError, "Size is not match in booleans: #{args}"
       end
       return take_by_array(vector) if vector.numeric?
-      return select_vars_by_keys(vector.to_a.map(&:to_sym)) if vector.string? || vector.type == :dictionary
+      return select_vars_by_keys(vector.to_a.map(&:to_sym)) if vector.string? || vector.dictionary?
 
       raise DataFrameArgumentError, "Invalid argument: #{args}"
     end
@@ -118,10 +118,10 @@ module RedAmber
         end
 
         normalized_indices = normalized_indices.floor.to_a.map(&:to_i) # round to integer array
-        return remove_all_values if normalized_indices == indices
+        return remove_all_values if normalized_indices == indices.to_a
         return self if normalized_indices.empty?
 
-        index_array = indices - normalized_indices
+        index_array = indices.to_a - normalized_indices
 
         datum = Arrow::Function.find(:take).execute([table, index_array])
         return DataFrame.new(datum.value)
@@ -168,14 +168,14 @@ module RedAmber
 
     # Undocumented
     # TODO: support for option {boundscheck: true}
-    def take(*indices)
-      indices.flatten!
-      return remove_all_values if indices.empty?
+    def take(*arg_indices)
+      arg_indices.flatten!
+      return remove_all_values if arg_indices.empty?
 
-      indices = indices[0] if indices.one? && !indices[0].is_a?(Numeric)
-      indices = Vector.new(indices) unless indices.is_a?(Vector)
+      arg_indices = arg_indices[0] if arg_indices.one? && !arg_indices[0].is_a?(Numeric)
+      arg_indices = Vector.new(arg_indices) unless arg_indices.is_a?(Vector)
 
-      take_by_array(indices)
+      take_by_array(arg_indices)
     end
 
     # Undocumented

data/lib/red_amber/data_frame_variable_operation.rb

@@ -23,12 +23,12 @@ module RedAmber
         elsif vec.numeric?
           key_vector.take(*vec).to_a
         elsif vec.string? || vec.dictionary?
-          picker
+          vec.to_a
         else
           raise DataFrameArgumentError, "Invalid argument #{args}"
         end
 
-      # DataFrame#[] creates a Vector with single key is specified.
+      # DataFrame#[] creates a Vector if single key is specified.
       # DataFrame#pick creates a DataFrame with single key.
       DataFrame.new(@table[ary])
     end
@@ -52,14 +52,14 @@ module RedAmber
         elsif vec.numeric?
           keys - key_vector.take(*vec).each.map(&:to_sym) # Array
         elsif vec.string? || vec.dictionary?
-          keys - dropper
+          keys - vec.to_a.map { _1&.to_sym } # Array
         else
           raise DataFrameArgumentError, "Invalid argument #{args}"
         end
 
       return DataFrame.new if ary.empty?
 
-      # DataFrame#[] creates a Vector with single key is specified.
+      # DataFrame#[] creates a Vector if single key is specified.
       # DataFrame#drop creates a DataFrame with single key.
       DataFrame.new(@table[ary])
     end

data/lib/red_amber/group.rb

@@ -3,35 +3,84 @@
 module RedAmber
   # group class
   class Group
+    include Enumerable # This feature is experimental
+
     # 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)
+      @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|
-        by(function, 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
 
-    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.new(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)
@@ -48,18 +97,50 @@ module RedAmber
 
     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] }
+    # @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

@@ -9,22 +9,10 @@ module RedAmber
       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))
@@ -51,7 +39,7 @@ module RedAmber
         elsif bg.nil? && en.nil?
           Array(0...vsize)
         else
-          Array[elem]
+          Array(elem)
         end
       when Enumerator
         elem.to_a

data/lib/red_amber/vector.rb

@@ -24,6 +24,8 @@ module RedAmber
             a
           in [Arrow::ChunkedArray => ca]
             ca
+          in [arrow_array_like] if arrow_array_like.respond_to?(:to_arrow_array)
+            arrow_array_like.to_arrow_array
           in [Range => r]
             Arrow::Array.new(Array(r))
           else
@@ -37,6 +39,11 @@ module RedAmber
     end
 
     attr_reader :data
+
+    def to_arrow_array
+      @data
+    end
+
     attr_accessor :key
 
     def to_s

data/lib/red_amber/vector_functions.rb

@@ -187,12 +187,6 @@ module RedAmber
     alias_method :ne, :not_equal
 
     def coerce(other)
-      case other
-      when Vector, Array, Arrow::Array
-        raise VectorArgumentError, "Size unmatch: #{size} != #{other.length}" unless size == other.length
-
-        [Vector.new(Array(other)), self]
-      end
       [Vector.new(Array(other) * size), self]
     end
 
@@ -271,8 +265,6 @@ module RedAmber
         find(function).execute([data, other.data], options)
       when Arrow::Array, Arrow::ChunkedArray, Arrow::Scalar, Array, Numeric, String, TrueClass, FalseClass
         find(function).execute([data, other], options)
-      else
-        raise VectorArgumentError, "Operand is not supported: #{other.class}"
       end
     end