red_amber 0.1.5 → 0.1.6

data/lib/red_amber/data_frame_selectable.rb

@@ -3,35 +3,94 @@
 module RedAmber
   # mix-in for the class DataFrame
   module DataFrameSelectable
-    # select columns: [symbol] or [string]
-    # select rows: [array of index], [range]
+    # select variables: [symbol] or [string]
+    # select observations: [array of index], [range]
     def [](*args)
+      args.flatten!
       raise DataFrameArgumentError, 'Empty dataframe' if empty?
-      raise DataFrameArgumentError, 'Empty argument' if args.empty?
-
-      if args.one?
-        case args[0]
-        when Vector
-          return select_obs_by_boolean(Arrow::BooleanArray.new(args[0].data))
-        when Arrow::BooleanArray
-          return select_obs_by_boolean(args[0])
-        when Array
-          return select_obs_by_boolean(Arrow::BooleanArray.new(args[0]))
-
-          # when Hash
-          # specify conditions to select by a Hash
-        end
+      return remove_all_values if args.empty? || args[0].nil?
+
+      vector = parse_to_vector(args)
+      if vector.boolean?
+        return filter_by_vector(vector.data) if vector.size == size
+
+        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
+
+      raise DataFrameArgumentError, "Invalid argument: #{args}"
+    end
+
+    # slice and select some observations to create sub DataFrame
+    def slice(*args, &block)
+      slicer = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+
+        slicer = instance_eval(&block)
+      end
+      slicer = [slicer].flatten
+
+      raise DataFrameArgumentError, 'Empty dataframe' if empty?
+      return remove_all_values if slicer.empty? || slicer[0].nil?
+
+      vector = parse_to_vector(slicer)
+      if vector.boolean?
+        return filter_by_vector(vector.data) if vector.size == size
+
+        raise DataFrameArgumentError, "Size is not match in booleans: #{slicer}"
       end
+      return take_by_array(vector) if vector.numeric?
+
+      raise DataFrameArgumentError, "Invalid argument #{slicer}"
+    end
+
+    # remove selected observations to create sub DataFrame
+    def remove(*args, &block)
+      remover = args
+      if block
+        raise DataFrameArgumentError, 'Must not specify both arguments and block.' unless args.empty?
+
+        remover = instance_eval(&block)
+      end
+      remover = [remover].flatten
+
+      raise DataFrameArgumentError, 'Empty dataframe' if empty?
+      return self if remover.empty? || remover[0].nil?
+
+      vector = parse_to_vector(remover)
+      if vector.boolean?
+        return filter_by_vector(vector.primitive_invert.data) if vector.size == size
+
+        raise DataFrameArgumentError, "Size is not match in booleans: #{remover}"
+      end
+      if vector.numeric?
+        raise DataFrameArgumentError, "Index out of range: #{vector.min}" if vector.min <= -size - 1
+
+        normalized_indices = (vector < 0).if_else(vector + size, vector) # normalize index from tail
+        if normalized_indices.max >= size
+          raise DataFrameArgumentError, "Index out of range: #{normalized_indices.max}"
+        end
 
-      return select_obs_by_boolean(args) if booleans?(args)
+        normalized_indices = normalized_indices.floor.to_a.map(&:to_i) # round to integer array
+        return remove_all_values if normalized_indices == indices
+        return self if normalized_indices.empty?
 
-      # expand Range like [1..3, 4] to [1, 2, 3, 4]
-      expanded = expand_range(args)
-      return map_indices(*expanded) if integers?(expanded)
-      return select_vars_by_keys(expanded.map(&:to_sym)) if sym_or_str?(expanded)
+        index_array = indices - normalized_indices
 
-      raise DataFrameArgumentError, "Invalid argument #{args}"
+        datum = Arrow::Function.find(:take).execute([table, index_array])
+        return DataFrame.new(datum.value)
+      end
+
+      raise DataFrameArgumentError, "Invalid argument #{remover}"
+    end
+
+    def remove_nil
+      func = Arrow::Function.find(:drop_null)
+      DataFrame.new(func.execute([table]).value)
     end
+    alias_method :drop_nil, :remove_nil
 
     # Select a variable by a key in String or Symbol
     def v(key)
@@ -43,24 +102,57 @@ module RedAmber
       variables[key.to_sym]
     end
 
-    def head(n_rows = 5)
-      raise DataFrameArgumentError, "Index is out of range #{n_rows}" if n_rows.negative?
+    def head(n_obs = 5)
+      raise DataFrameArgumentError, "Index is out of range #{n_obs}" if n_obs.negative?
 
-      self[0...[n_rows, size].min]
+      self[0...[n_obs, size].min]
     end
 
-    def tail(n_rows = 5)
-      raise DataFrameArgumentError, "Index is out of range #{n_rows}" if n_rows.negative?
+    def tail(n_obs = 5)
+      raise DataFrameArgumentError, "Index is out of range #{n_obs}" if n_obs.negative?
 
-      self[-[n_rows, size].min..]
+      self[-[n_obs, size].min..]
     end
 
-    def first(n_rows = 1)
-      head(n_rows)
+    def first(n_obs = 1)
+      head(n_obs)
     end
 
-    def last(n_rows = 1)
-      tail(n_rows)
+    def last(n_obs = 1)
+      tail(n_obs)
+    end
+
+    # Undocumented
+    # TODO: support for option {boundscheck: true}
+    def take(*indices)
+      indices.flatten!
+      return remove_all_values if indices.empty?
+
+      indices = indices[0] if indices.one? && !indices[0].is_a?(Numeric)
+      indices = Vector.new(indices) unless indices.is_a?(Vector)
+
+      take_by_array(indices)
+    end
+
+    # Undocumented
+    # TODO: support for option {null_selection_behavior: :drop}
+    def filter(*booleans)
+      booleans.flatten!
+      return remove_all_values if booleans.empty?
+
+      b = booleans[0]
+      case b
+      when Vector
+        raise DataFrameArgumentError, 'Argument is not a boolean.' unless b.boolean?
+
+        filter_by_vector(b.data)
+      when Arrow::BooleanArray
+        filter_by_vector(b)
+      else
+        raise DataFrameArgumentError, 'Argument is not a boolean.' unless booleans?(booleans)
+
+        filter_by_vector(Arrow::BooleanArray.new(booleans))
+      end
     end
 
     private
@@ -75,5 +167,32 @@ module RedAmber
         DataFrame.new(@table[keys])
       end
     end
+
+    # Accepts indices by numeric Vector
+    def take_by_array(indices)
+      raise DataFrameArgumentError, "Indices must be a numeric Vector: #{indices}" unless indices.numeric?
+      raise DataFrameArgumentError, "Index out of range: #{indices.min}" if indices.min <= -size - 1
+
+      normalized_indices = (indices < 0).if_else(indices + size, indices) # normalize index from tail
+      raise DataFrameArgumentError, "Index out of range: #{normalized_indices.max}" if normalized_indices.max >= size
+
+      index_array = Arrow::UInt64ArrayBuilder.build(normalized_indices.data) # round to integer array
+
+      datum = Arrow::Function.find(:take).execute([table, index_array])
+      DataFrame.new(datum.value)
+    end
+
+    # Accepts booleans by Arrow::BooleanArray
+    def filter_by_vector(boolean_array)
+      raise DataFrameArgumentError, 'Booleans must be same size as self.' unless boolean_array.length == size
+
+      datum = Arrow::Function.find(:filter).execute([table, boolean_array])
+      DataFrame.new(datum.value)
+    end
+
+    # return a DataFrame with same keys as self without values
+    def remove_all_values
+      filter_by_vector(Arrow::BooleanArray.new([false] * size))
+    end
   end
 end

data/lib/red_amber/data_frame_variable_operation.rb

@@ -129,5 +129,9 @@ module RedAmber
         arrays << Arrow::ChunkedArray.new([a])
       end
     end
+
+    def keys_by_booleans(booleans)
+      keys.select.with_index { |_, i| booleans[i] }
+    end
   end
 end

data/lib/red_amber/helper.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module RedAmber
+  # mix-in for the class DataFrame
+  module Helper
+    private
+
+    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 sym_or_str?(enum)
+      enum.all? { |e| e.is_a?(Symbol) || e.is_a?(String) }
+    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)
+      a = args.reduce([]) do |accum, elem|
+        accum.concat(normalize_element(elem))
+      end
+      Vector.new(a)
+    end
+
+    def normalize_element(elem)
+      case elem
+      when Numeric, String, Symbol, TrueClass, FalseClass, NilClass
+        [elem]
+      when Range
+        both_end = [elem.begin, elem.end]
+        both_end[1] -= 1 if elem.exclude_end? && elem.end.is_a?(Integer)
+
+        if both_end.any?(Integer) || both_end.all?(&:nil?)
+          if both_end.any? { |e| e&.>=(size) || e&.<(-size) }
+            raise DataFrameArgumentError, "Index out of range: #{elem} for 0..#{size - 1}"
+          end
+
+          (0...size).to_a[elem]
+        else
+          elem.to_a
+        end
+      else
+        Array(elem)
+      end
+    end
+  end
+end

data/lib/red_amber/vector.rb

@@ -1,25 +1,37 @@
 # frozen_string_literal: true
 
 module RedAmber
-  # Columnar data object
+  # Values in variable (columnar) data object
   #   @data : holds Arrow::ChunkedArray
   class Vector
     # mix-in
-    include VectorCompensable
     include VectorFunctions
+    include VectorUpdatable
+    include VectorSelectable
+    include Helper
 
-    # chunked_array may come from column.data
-    def initialize(array)
+    def initialize(*array)
       @key = nil # default is 'headless'
-      case array
-      when Vector
-        @data = array.data
-      when Arrow::Array, Arrow::ChunkedArray
-        @data = array
-      when Array
-        @data = Arrow::Array.new(array)
+      if array.empty? || array[0].nil?
+        Vector.new([])
       else
-        raise VectorArgumentError, 'Unknown array in argument'
+        array.flatten!
+        case array[0]
+        when Vector
+          @data = array[0].data
+          return
+        when Arrow::Array, Arrow::ChunkedArray
+          @data = array[0]
+          return
+        when Range
+          @data = Arrow::Array.new(Array(array[0]))
+          return
+        end
+        begin
+          @data = Arrow::Array.new(Array(array))
+        rescue Error
+          raise VectorArgumentError, "Invalid argument: #{array}"
+        end
       end
     end
 
@@ -52,6 +64,16 @@ module RedAmber
     alias_method :to_a, :values
     alias_method :entries, :values
 
+    def indices
+      (0...size).to_a
+    end
+    alias_method :indexes, :indices
+    alias_method :indeces, :indices
+
+    def to_ary
+      to_a
+    end
+
     def size
       # only defined :length in Arrow?
       @data.length
@@ -60,6 +82,10 @@ module RedAmber
     alias_method :n_rows, :size
     alias_method :nrow, :size
 
+    def empty?
+      size.zero?
+    end
+
     def type
       @data.value_type.nick.to_sym
     end
@@ -124,5 +150,9 @@ module RedAmber
     def n_nans
       numeric? ? is_nan.to_a.count(true) : 0
     end
+
+    def has_nil?
+      is_nil.any
+    end
   end
 end

data/lib/red_amber/vector_functions.rb

@@ -16,11 +16,13 @@ module RedAmber
     unary_aggregations.each do |function|
       define_method(function) do |opts: nil|
         datum = exec_func_unary(function, options: opts)
-        take_out_scalar(datum)
+        get_scalar(datum)
       end
     end
     alias_method :median, :approximate_median
     alias_method :count_uniq, :count_distinct
+    alias_method :all?, :all
+    alias_method :any?, :any
 
     def unbiased_variance
       variance(opts: { ddof: 1 })
@@ -47,7 +49,7 @@ module RedAmber
     unary_element_wise.each do |function|
       define_method(function) do |opts: nil|
         datum = exec_func_unary(function, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
     end
     alias_method :is_nil, :is_null
@@ -72,12 +74,12 @@ module RedAmber
     unary_element_wise_op.each do |function, operator|
       define_method(function) do |opts: nil|
         datum = exec_func_unary(function, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
 
       define_method(operator) do |opts: nil|
         datum = exec_func_unary(function, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
     end
     alias_method :not, :invert
@@ -95,7 +97,7 @@ module RedAmber
     binary_element_wise.each do |function|
       define_method(function) do |other, opts: nil|
         datum = exec_func_binary(function, other, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
     end
 
@@ -111,7 +113,7 @@ module RedAmber
     logical_binary_element_wise.each do |method, function|
       define_method(method) do |other, opts: nil|
         datum = exec_func_binary(function, other, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
     end
 
@@ -144,12 +146,12 @@ module RedAmber
     binary_element_wise_op.each do |function, operator|
       define_method(function) do |other, opts: nil|
         datum = exec_func_binary(function, other, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
 
       define_method(operator) do |other, opts: nil|
         datum = exec_func_binary(function, other, options: opts)
-        take_out_element_wise(datum)
+        Vector.new(datum.value)
       end
     end
     alias_method :eq, :equal
@@ -159,8 +161,17 @@ module RedAmber
     alias_method :lt, :less
     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
+
     # (array functions)
-    # array_filter, array_take
     # dictionary_encode,
     # partition_nth_indices,
     # quarter, quarters_between,
@@ -192,17 +203,17 @@ module RedAmber
     # strptime, subsecond, us_week, week, weeks_between, year, year_month_day, years_between
 
     # (onditional)
-    # case_when, cast, if_else
+    # case_when, cast,
 
     # (indices)
     # choose, index_in, index_in_meta_binary, indices_nonzero
 
     # (others)
-    # coalesce, drop_null,
-    # filter, is_in, is_in_meta_binary,
+    # coalesce,
+    # is_in_meta_binary,
     # list_element, list_flatten, list_parent_indices, list_value_length, make_struct,
     # max_element_wise, min_element_wise, random, select_k_unstable,
-    # sort_indices, struct_field, take
+    # struct_field,
 
     private # =======
 
@@ -221,7 +232,7 @@ module RedAmber
       end
     end
 
-    def take_out_scalar(datum)
+    def get_scalar(datum)
       output = datum.value
       case output
       when Arrow::StringScalar then output.to_s
@@ -232,10 +243,6 @@ module RedAmber
       end
     end
 
-    def take_out_element_wise(datum)
-      Vector.new(datum.value)
-    end
-
     module_function # ======
 
     def find(function_name)

data/lib/red_amber/vector_selectable.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+# Available functions in Arrow are shown by `Arrow::Function.all.map(&:name)`
+# reference: https://arrow.apache.org/docs/cpp/compute.html
+
+module RedAmber
+  # mix-ins for class Vector
+  # Functions to select some data.
+  module VectorSelectable
+    def drop_nil
+      datum = find(:drop_null).execute([data])
+      Vector.new(datum.value)
+    end
+
+    # vector calculation version of selection by indices
+    # TODO: support for option {boundscheck: true}
+    def take(*indices)
+      indices.flatten!
+      return Vector.new([]) if indices.empty?
+
+      indices = indices[0] if indices.one? && !indices[0].is_a?(Numeric)
+      indices = Vector.new(indices) unless indices.is_a?(Vector)
+
+      take_by_vector(indices) # returns sub Vector
+    end
+
+    # TODO: support for option {null_selection_behavior: :drop}
+    def filter(*booleans)
+      booleans.flatten!
+      return Vector.new([]) if booleans.empty?
+
+      b = booleans[0]
+      boolean_array =
+        case b
+        when Vector
+          raise VectorTypeError, 'Argument is not a boolean.' unless b.boolean?
+
+          b.data
+        when Arrow::BooleanArray
+          b
+        else
+          raise VectorTypeError, 'Argument is not a boolean.' unless booleans?(booleans)
+
+          Arrow::BooleanArray.new(booleans)
+        end
+
+      filter_by_array(boolean_array) # returns sub Vector
+    end
+
+    #   @param indices
+    #   @param booleans
+    def [](*args)
+      args.flatten!
+      return Vector.new([]) if args.empty?
+
+      arg = args[0]
+      case arg
+      when Vector
+        return take_by_vector(arg) if arg.numeric?
+        return filter_by_array(arg.data) if arg.boolean?
+
+        raise VectorTypeError, "Argument must be numeric or boolean: #{arg}"
+      when Arrow::BooleanArray
+        return filter_by_array(arg)
+      when Arrow::Array
+        array = arg
+      else
+        unless arg.is_a?(Numeric) || booleans?([arg])
+          raise VectorArgumentError, "Argument must be numeric or boolean: #{args}"
+        end
+      end
+      array ||= Arrow::Array.new(args)
+      return filter_by_array(array) if array.is_a?(Arrow::BooleanArray)
+
+      vector = Vector.new(array)
+      return take_by_vector(vector) if vector.numeric?
+
+      raise VectorArgumentError, "Invalid argument: #{args}"
+    end
+
+    #   @param values [Array, Arrow::Array, Vector]
+    def is_in(*values)
+      values.flatten!
+      array =
+        case values[0]
+        when Vector
+          values[0].data
+        when Arrow::Array
+          values[0]
+        end
+      array ||= data.class.new(values)
+      Vector.new(data.is_in(array))
+    end
+
+    # Arrow's support required
+    def index(element)
+      to_a.index(element)
+    end
+
+    private
+
+    # Accepts indices by numeric Vector
+    def take_by_vector(indices)
+      raise VectorTypeError, "Indices must be numeric Vector: #{indices}" unless indices.numeric?
+      raise VectorArgumentError, "Index out of range: #{indices.min}" if indices.min <= -size - 1
+
+      normalized_indices = (indices < 0).if_else(indices + size, indices) # normalize index from tail
+      raise VectorArgumentError, "Index out of range: #{normalized_indices.max}" if normalized_indices.max >= size
+
+      index_array = Arrow::UInt64ArrayBuilder.build(normalized_indices.data) # round to integer array
+
+      datum = find(:array_take).execute([data, index_array])
+      Vector.new(datum.value)
+    end
+
+    # Accepts booleans by Arrow::BooleanArray
+    def filter_by_array(boolean_array)
+      raise VectorArgumentError, 'Booleans must be same size as self.' unless boolean_array.length == size
+
+      datum = find(:array_filter).execute([data, boolean_array])
+      Vector.new(datum.value)
+    end
+  end
+end