polars-df 0.1.1 → 0.1.3

data/lib/polars/lazy_functions.rb

@@ -1,14 +1,87 @@
 module Polars
   module LazyFunctions
+    # Return an expression representing a column in a DataFrame.
+    #
+    # @return [Expr]
     def col(name)
-      name = name.to_s if name.is_a?(Symbol)
-      Utils.wrap_expr(RbExpr.col(name))
+      if name.is_a?(Series)
+        name = name.to_a
+      end
+
+      if name.is_a?(Array)
+        if name.length == 0 || name[0].is_a?(String) || name[0].is_a?(Symbol)
+          name = name.map { |v| v.is_a?(Symbol) ? v.to_s : v }
+          Utils.wrap_expr(RbExpr.cols(name))
+        elsif Utils.is_polars_dtype(name[0])
+          raise Todo
+          # Utils.wrap_expr(_dtype_cols(name))
+        else
+          raise ArgumentError, "Expected list values to be all `str` or all `DataType`"
+        end
+      else
+        name = name.to_s if name.is_a?(Symbol)
+        Utils.wrap_expr(RbExpr.col(name))
+      end
     end
 
+    # Alias for an element in evaluated in an `eval` expression.
+    #
+    # @return [Expr]
+    #
+    # @example A horizontal rank computation by taking the elements of a list
+    #   df = Polars::DataFrame.new({"a" => [1, 8, 3], "b" => [4, 5, 2]})
+    #   df.with_column(
+    #     Polars.concat_list(["a", "b"]).arr.eval(Polars.element.rank).alias("rank")
+    #   )
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬────────────┐
+    #   # │ a   ┆ b   ┆ rank       │
+    #   # │ --- ┆ --- ┆ ---        │
+    #   # │ i64 ┆ i64 ┆ list[f32]  │
+    #   # ╞═════╪═════╪════════════╡
+    #   # │ 1   ┆ 4   ┆ [1.0, 2.0] │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 8   ┆ 5   ┆ [2.0, 1.0] │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 3   ┆ 2   ┆ [2.0, 1.0] │
+    #   # └─────┴─────┴────────────┘
     def element
       col("")
     end
 
+    # Count the number of values in this column/context.
+    #
+    # @param column [String, Series, nil]
+    #     If dtype is:
+    #
+    #     * `Series` : count the values in the series.
+    #     * `String` : count the values in this column.
+    #     * `None` : count the number of values in this context.
+    #
+    # @return [Expr, Integer]
+    def count(column = nil)
+      if column.nil?
+        return Utils.wrap_expr(RbExpr.count)
+      end
+
+      if column.is_a?(Series)
+        column.len
+      else
+        col(column).count
+      end
+    end
+
+    # Aggregate to list.
+    #
+    # @return [Expr]
+    def to_list(name)
+      col(name).list
+    end
+
+    # Get the standard deviation.
+    #
+    # @return [Object]
     def std(column, ddof: 1)
       if column.is_a?(Series)
         column.std(ddof: ddof)
@@ -17,6 +90,9 @@ module Polars
       end
     end
 
+    # Get the variance.
+    #
+    # @return [Object]
     def var(column, ddof: 1)
       if column.is_a?(Series)
         column.var(ddof: ddof)
@@ -25,6 +101,16 @@ module Polars
       end
     end
 
+    # Get the maximum value.
+    #
+    # @param column [Object]
+    #   Column(s) to be used in aggregation. Will lead to different behavior based on
+    #   the input:
+    #
+    #   - [String, Series] -> aggregate the maximum value of that column.
+    #   - [Array<Expr>] -> aggregate the maximum value horizontally.
+    #
+    # @return [Expr, Object]
     def max(column)
       if column.is_a?(Series)
         column.max
@@ -37,6 +123,16 @@ module Polars
       end
     end
 
+    # Get the minimum value.
+    #
+    # @param column [Object]
+    #   Column(s) to be used in aggregation. Will lead to different behavior based on
+    #   the input:
+    #
+    #   - [String, Series] -> aggregate the minimum value of that column.
+    #   - [Array<Expr>] -> aggregate the minimum value horizontally.
+    #
+    # @return [Expr, Object]
     def min(column)
       if column.is_a?(Series)
         column.min
@@ -49,6 +145,9 @@ module Polars
       end
     end
 
+    # Sum values in a column/Series, or horizontally across list of columns/expressions.
+    #
+    # @return [Object]
     def sum(column)
       if column.is_a?(Series)
         column.sum
@@ -59,10 +158,13 @@ module Polars
         # TODO
         Utils.wrap_expr(_sum_exprs(exprs))
       else
-        raise "todo"
+        raise Todo
       end
     end
 
+    # Get the mean value.
+    #
+    # @return [Expr, Float]
     def mean(column)
       if column.is_a?(Series)
         column.mean
@@ -71,10 +173,16 @@ module Polars
       end
     end
 
+    # Get the mean value.
+    #
+    # @return [Expr, Float]
     def avg(column)
       mean(column)
     end
 
+    # Get the median value.
+    #
+    # @return [Object]
     def median(column)
       if column.is_a?(Series)
         column.median
@@ -83,10 +191,160 @@ module Polars
       end
     end
 
+    # def n_unique
+    # end
+
+    # Get the first value.
+    #
+    # @return [Object]
+    def first(column = nil)
+      if column.nil?
+        return Utils.wrap_expr(RbExpr.first)
+      end
+
+      if column.is_a?(Series)
+        if column.len > 0
+          column[0]
+        else
+          raise IndexError, "The series is empty, so no first value can be returned."
+        end
+      else
+        col(column).first
+      end
+    end
+
+    # def last
+    # end
+
+    # def head
+    # end
+
+    # def tail
+    # end
+
+    # Return an expression representing a literal value.
+    #
+    # @return [Expr]
     def lit(value)
+      if value.is_a?(Polars::Series)
+        name = value.name
+        value = value._s
+        e = Utils.wrap_expr(RbExpr.lit(value))
+        if name == ""
+          return e
+        end
+        return e.alias(name)
+      end
+
       Utils.wrap_expr(RbExpr.lit(value))
     end
 
+    # def cumsum
+    # end
+
+    # def spearman_rank_corr
+    # end
+
+    # def pearson_corr
+    # end
+
+    # def cov
+    # end
+
+    # def map
+    # end
+
+    # def apply
+    # end
+
+    # Accumulate over multiple columns horizontally/ row wise with a left fold.
+    #
+    # @return [Expr]
+    def fold(acc, f, exprs)
+      acc = Utils.expr_to_lit_or_expr(acc, str_to_lit: true)
+      if exprs.is_a?(Expr)
+        exprs = [exprs]
+      end
+
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(RbExpr.fold(acc._rbexpr, f, exprs))
+    end
+
+    # def reduce
+    # end
+
+    # def cumfold
+    # end
+
+    # def cumreduce
+    # end
+
+    # def any
+    # end
+
+    # def exclude
+    # end
+
+    # Do one of two things.
+    #
+    # * function can do a columnwise or elementwise AND operation
+    # * a wildcard column selection
+    #
+    # @param name [Object]
+    #   If given this function will apply a bitwise & on the columns.
+    #
+    # @return [Expr]
+    #
+    # @example Sum all columns
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [1, 2, 3], "b" => ["hello", "foo", "bar"], "c" => [1, 1, 1]}
+    #   )
+    #   df.select(Polars.all.sum)
+    #   # =>
+    #   # shape: (1, 3)
+    #   # ┌─────┬──────┬─────┐
+    #   # │ a   ┆ b    ┆ c   │
+    #   # │ --- ┆ ---  ┆ --- │
+    #   # │ i64 ┆ str  ┆ i64 │
+    #   # ╞═════╪══════╪═════╡
+    #   # │ 6   ┆ null ┆ 3   │
+    #   # └─────┴──────┴─────┘
+    def all(name = nil)
+      if name.nil?
+        col("*")
+      elsif name.is_a?(String) || name.is_a?(Symbol)
+        col(name).all
+      else
+        raise Todo
+      end
+    end
+
+    # def groups
+    # end
+
+    # def quantile
+    # end
+
+    # Create a range expression (or Series).
+    #
+    # This can be used in a `select`, `with_column`, etc. Be sure that the resulting
+    # range size is equal to the length of the DataFrame you are collecting.
+    #
+    # @param low [Integer, Expr, Series]
+    #   Lower bound of range.
+    # @param high [Integer, Expr, Series]
+    #   Upper bound of range.
+    # @param step [Integer]
+    #   Step size of the range.
+    # @param eager [Boolean]
+    #   If eager evaluation is `True`, a Series is returned instead of an Expr.
+    # @param dtype [Symbol]
+    #   Apply an explicit integer dtype to the resulting expression (default is Int64).
+    #
+    # @return [Expr, Series]
+    #
+    # @example
+    #   df.lazy.filter(Polars.col("foo") < Polars.arange(0, 100)).collect
     def arange(low, high, step: 1, eager: false, dtype: nil)
       low = Utils.expr_to_lit_or_expr(low, str_to_lit: false)
       high = Utils.expr_to_lit_or_expr(high, str_to_lit: false)
@@ -106,16 +364,179 @@ module Polars
       end
     end
 
-    def all(name = nil)
-      if name.nil?
-        col("*")
-      elsif name.is_a?(String) || name.is_a?(Symbol)
-        col(name).all
+    # def argsort_by
+    # end
+
+    # def duration
+    # end
+
+    # def format
+    # end
+
+    # Concat the arrays in a Series dtype List in linear time.
+    #
+    # @return [Expr]
+    def concat_list(exprs)
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(RbExpr.concat_lst(exprs))
+    end
+
+    # def collect_all
+    # end
+
+    # Run polars expressions without a context.
+    #
+    # @return [DataFrame]
+    def select(exprs)
+      DataFrame.new([]).select(exprs)
+    end
+
+    # Collect several columns into a Series of dtype Struct.
+    #
+    # @param exprs [Object]
+    #   Columns/Expressions to collect into a Struct
+    # @param eager [Boolean]
+    #   Evaluate immediately
+    #
+    # @return [Object]
+    #
+    # @example
+    #   Polars::DataFrame.new(
+    #     {
+    #       "int" => [1, 2],
+    #       "str" => ["a", "b"],
+    #       "bool" => [true, nil],
+    #       "list" => [[1, 2], [3]],
+    #     }
+    #   ).select([Polars.struct(Polars.all()).alias("my_struct")])
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────────────────────┐
+    #   # │ my_struct           │
+    #   # │ ---                 │
+    #   # │ struct[4]           │
+    #   # ╞═════════════════════╡
+    #   # │ {1,"a",true,[1, 2]} │
+    #   # ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ {2,"b",null,[3]}    │
+    #   # └─────────────────────┘
+    #
+    # @example Only collect specific columns as a struct:
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [1, 2, 3, 4], "b" => ["one", "two", "three", "four"], "c" => [9, 8, 7, 6]}
+    #   )
+    #   df.with_column(pl.struct(pl.col(["a", "b"])).alias("a_and_b"))
+    #   # =>
+    #   # shape: (4, 4)
+    #   # ┌─────┬───────┬─────┬─────────────┐
+    #   # │ a   ┆ b     ┆ c   ┆ a_and_b     │
+    #   # │ --- ┆ ---   ┆ --- ┆ ---         │
+    #   # │ i64 ┆ str   ┆ i64 ┆ struct[2]   │
+    #   # ╞═════╪═══════╪═════╪═════════════╡
+    #   # │ 1   ┆ one   ┆ 9   ┆ {1,"one"}   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 2   ┆ two   ┆ 8   ┆ {2,"two"}   │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 3   ┆ three ┆ 7   ┆ {3,"three"} │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+    #   # │ 4   ┆ four  ┆ 6   ┆ {4,"four"}  │
+    #   # └─────┴───────┴─────┴─────────────┘
+    def struct(exprs, eager: false)
+      if eager
+        Polars.select(struct(exprs, eager: false)).to_series
+      end
+      exprs = Utils.selection_to_rbexpr_list(exprs)
+      Utils.wrap_expr(_as_struct(exprs))
+    end
+
+    # Repeat a single value n times.
+    #
+    # @param value [Object]
+    #   Value to repeat.
+    # @param n [Integer]
+    #   Repeat `n` times.
+    # @param eager [Boolean]
+    #   Run eagerly and collect into a `Series`.
+    # @param name [String]
+    #   Only used in `eager` mode. As expression, use `alias`.
+    #
+    # @return [Expr]
+    def repeat(value, n, eager: false, name: nil)
+      if eager
+        if name.nil?
+          name = ""
+        end
+        dtype = py_type_to_dtype(type(value))
+        Series._repeat(name, value, n, dtype)
       else
-        raise "todo"
+        if n.is_a?(Integer)
+          n = lit(n)
+        end
+        Utils.wrap_expr(RbExpr.repeat(value, n._rbexpr))
       end
     end
 
+    # Return indices where `condition` evaluates `true`.
+    #
+    # @param condition [Expr]
+    #   Boolean expression to evaluate
+    # @param eager [Boolean]
+    #   Whether to apply this function eagerly (as opposed to lazily).
+    #
+    # @return [Expr, Series]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [1, 2, 3, 4, 5]})
+    #   df.select(
+    #     [
+    #       Polars.arg_where(Polars.col("a") % 2 == 0)
+    #     ]
+    #   ).to_series
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         1
+    #   #         3
+    #   # ]
+    def arg_where(condition, eager: false)
+      if eager
+        if !condition.is_a?(Series)
+          raise ArgumentError, "expected 'Series' in 'arg_where' if 'eager=True', got #{condition.class.name}"
+        end
+        condition.to_frame.select(arg_where(Polars.col(condition.name))).to_series
+      else
+        condition = Utils.expr_to_lit_or_expr(condition, str_to_lit: true)
+        Utils.wrap_expr(_arg_where(condition._rbexpr))
+      end
+    end
+
+    # def coalesce
+    # end
+
+    # def from_epoch
+    # end
+
+    # Start a "when, then, otherwise" expression.
+    #
+    # @return [When]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"foo" => [1, 3, 4], "bar" => [3, 4, 0]})
+    #   df.with_column(Polars.when(Polars.col("foo") > 2).then(Polars.lit(1)).otherwise(Polars.lit(-1)))
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌─────┬─────┬─────────┐
+    #   # │ foo ┆ bar ┆ literal │
+    #   # │ --- ┆ --- ┆ ---     │
+    #   # │ i64 ┆ i64 ┆ i32     │
+    #   # ╞═════╪═════╪═════════╡
+    #   # │ 1   ┆ 3   ┆ -1      │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
+    #   # │ 3   ┆ 4   ┆ 1       │
+    #   # ├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
+    #   # │ 4   ┆ 0   ┆ 1       │
+    #   # └─────┴─────┴─────────┘
     def when(expr)
       expr = Utils.expr_to_lit_or_expr(expr)
       pw = RbExpr.when(expr._rbexpr)

data/lib/polars/lazy_group_by.rb

@@ -1,13 +1,92 @@
 module Polars
   class LazyGroupBy
+    # @private
     def initialize(lgb, lazyframe_class)
       @lgb = lgb
       @lazyframe_class = lazyframe_class
     end
 
+    # Describe the aggregation that need to be done on a group.
+    #
+    # @return [LazyFrame]
     def agg(aggs)
       rbexprs = Utils.selection_to_rbexpr_list(aggs)
       @lazyframe_class._from_rbldf(@lgb.agg(rbexprs))
     end
+
+    # Get the first `n` rows of each group.
+    #
+    # @param n [Integer]
+    #   Number of rows to return.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "letters" => ["c", "c", "a", "c", "a", "b"],
+    #       "nrs" => [1, 2, 3, 4, 5, 6]
+    #     }
+    #   )
+    #   df.groupby("letters").head(2).sort("letters")
+    #   # =>
+    #   # shape: (5, 2)
+    #   # ┌─────────┬─────┐
+    #   # │ letters ┆ nrs │
+    #   # │ ---     ┆ --- │
+    #   # │ str     ┆ i64 │
+    #   # ╞═════════╪═════╡
+    #   # │ a       ┆ 3   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ a       ┆ 5   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ b       ┆ 6   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ c       ┆ 1   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ c       ┆ 2   │
+    #   # └─────────┴─────┘
+    def head(n = 5)
+      @lazyframe_class._from_rbldf(@lgb.head(n))
+    end
+
+    # Get the last `n` rows of each group.
+    #
+    # @param n [Integer]
+    #   Number of rows to return.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "letters" => ["c", "c", "a", "c", "a", "b"],
+    #       "nrs" => [1, 2, 3, 4, 5, 6]
+    #     }
+    #   )
+    #   df.groupby("letters").tail(2).sort("letters")
+    #   # =>
+    #   # shape: (5, 2)
+    #   # ┌─────────┬─────┐
+    #   # │ letters ┆ nrs │
+    #   # │ ---     ┆ --- │
+    #   # │ str     ┆ i64 │
+    #   # ╞═════════╪═════╡
+    #   # │ a       ┆ 3   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ a       ┆ 5   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ b       ┆ 6   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ c       ┆ 2   │
+    #   # ├╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌┤
+    #   # │ c       ┆ 4   │
+    #   # └─────────┴─────┘
+    def tail(n = 5)
+      @lazyframe_class._from_rbldf(@lgb.tail(n))
+    end
+
+    # def apply
+    # end
   end
 end

data/lib/polars/list_expr.rb

@@ -1,7 +1,10 @@
 module Polars
+  # Namespace for list related expressions.
   class ListExpr
+    # @private
     attr_accessor :_rbexpr
 
+    # @private
     def initialize(expr)
       self._rbexpr = expr._rbexpr
     end
@@ -41,6 +44,7 @@ module Polars
     # def concat
     # end
 
+    #
     def get(index)
       index = Utils.expr_to_lit_or_expr(index, str_to_lit: false)._rbexpr
       Utils.wrap_expr(_rbexpr.lst_get(index))
@@ -101,6 +105,7 @@ module Polars
     #   Utils.wrap_expr(_rbexpr.lst_to_struct(n_field_strategy, name_generator))
     # end
 
+    #
     def eval(expr, parallel: false)
        Utils.wrap_expr(_rbexpr.lst_eval(expr._rbexpr, parallel))
     end

data/lib/polars/meta_expr.rb

@@ -1,31 +1,52 @@
 module Polars
+  # Namespace for expressions on a meta level.
   class MetaExpr
+    # @private
     attr_accessor :_rbexpr
 
+    # @private
     def initialize(expr)
       self._rbexpr = expr._rbexpr
     end
 
+    # Equal.
+    #
+    # @return [Boolean]
     def ==(other)
       _rbexpr.meta_eq(other._rbexpr)
     end
 
+    # Not equal.
+    #
+    # @return [Boolean]
     def !=(other)
       !(self == other)
     end
 
+    # Pop the latest expression and return the input(s) of the popped expression.
+    #
+    # @return [Array]
     def pop
       _rbexpr.meta_pop.map { |e| Utils.wrap_expr(e) }
     end
 
+    # Get a list with the root column name.
+    #
+    # @return [Array]
     def root_names
       _rbexpr.meta_roots
     end
 
+    # Get the column name that this expression would produce.
+    #
+    # @return [String]
     def output_name
       _rbexpr.meta_output_name
     end
 
+    # Undo any renaming operation like `alias` or `keep_name`.
+    #
+    # @return [Expr]
     def undo_aliases
       Utils.wrap_expr(_rbexpr.meta_undo_aliases)
     end