polars-df 0.15.0 → 0.26.0

data/lib/polars/lazy_group_by.rb

@@ -6,6 +6,45 @@ module Polars
       @lgb = lgb
     end
 
+    # Filter groups with a list of predicates after aggregation.
+    #
+    # Using this method is equivalent to adding the predicates to the aggregation and
+    # filtering afterwards.
+    #
+    # This method can be chained and all conditions will be combined using `&`.
+    #
+    # @param predicates [Array]
+    #   Expressions that evaluate to a boolean value for each group. Typically, this
+    #   requires the use of an aggregation function. Multiple predicates are
+    #   combined using `&`.
+    #
+    # @return [LazyGroupBy]
+    #
+    # @example Only keep groups that contain more than one element.
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => ["a", "b", "a", "b", "c"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("a").having(
+    #     Polars.len > 1
+    #   ).agg.collect
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ str │
+    #   # ╞═════╡
+    #   # │ b   │
+    #   # │ a   │
+    #   # └─────┘
+    def having(*predicates)
+      rbexprs = Utils.parse_into_list_of_expressions(*predicates)
+      @lgb = @lgb.having(rbexprs)
+      self
+    end
+
     # Compute aggregations for each group of a group by operation.
     #
     # @param aggs [Array]
@@ -111,6 +150,65 @@ module Polars
       Utils.wrap_ldf(@lgb.agg(rbexprs))
     end
 
+    # Apply a custom/user-defined function (UDF) over the groups as a new DataFrame.
+    #
+    # @note
+    #   This method is much slower than the native expressions API.
+    #   Only use it if you cannot implement your logic otherwise.
+    #
+    # Using this is considered an anti-pattern as it will be very slow because:
+    #
+    # - it forces the engine to materialize the whole `DataFrames` for the groups.
+    # - it is not parallelized
+    # - it blocks optimizations as the passed python function is opaque to the
+    #   optimizer
+    #
+    # The idiomatic way to apply custom functions over multiple columns is using:
+    #
+    # `Polars.struct([my_columns]).apply { |struct_series| ... }`
+    #
+    # @param schema [Object]
+    #   Schema of the output function. This has to be known statically. If the
+    #   given schema is incorrect, this is a bug in the caller's query and may
+    #   lead to errors. If set to None, polars assumes the schema is unchanged.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "id" => [0, 1, 2, 3, 4],
+    #       "color" => ["red", "green", "green", "red", "red"],
+    #       "shape" => ["square", "triangle", "square", "triangle", "square"]
+    #     }
+    #   )
+    #   (
+    #     df.lazy
+    #       .group_by("color")
+    #       .map_groups(nil) { |group_df| group_df.sample(n: 2) }
+    #       .collect
+    #   )
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌─────┬───────┬──────────┐
+    #   # │ id  ┆ color ┆ shape    │
+    #   # │ --- ┆ ---   ┆ ---      │
+    #   # │ i64 ┆ str   ┆ str      │
+    #   # ╞═════╪═══════╪══════════╡
+    #   # │ 1   ┆ green ┆ triangle │
+    #   # │ 2   ┆ green ┆ square   │
+    #   # │ 4   ┆ red   ┆ square   │
+    #   # │ 3   ┆ red   ┆ triangle │
+    #   # └─────┴───────┴──────────┘
+    def map_groups(
+      schema,
+      &function
+    )
+      Utils.wrap_ldf(
+        @lgb.map_groups(->(df) { function.(Utils.wrap_df(df))._df }, schema)
+      )
+    end
+
     # Get the first `n` rows of each group.
     #
     # @param n [Integer]
@@ -175,7 +273,343 @@ module Polars
       Utils.wrap_ldf(@lgb.tail(n))
     end
 
-    # def apply
-    # end
+    # Aggregate the groups into Series.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => ["one", "two", "one", "two"],
+    #       "b" => [1, 2, 3, 4]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("a", maintain_order: true).all.collect
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬───────────┐
+    #   # │ a   ┆ b         │
+    #   # │ --- ┆ ---       │
+    #   # │ str ┆ list[i64] │
+    #   # ╞═════╪═══════════╡
+    #   # │ one ┆ [1, 3]    │
+    #   # │ two ┆ [2, 4]    │
+    #   # └─────┴───────────┘
+    def all
+      agg(F.all)
+    end
+
+    # Return the number of rows in each group.
+    #
+    # @param name [String]
+    #   Assign a name to the resulting column; if unset, defaults to "len".
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   lf = Polars::LazyFrame.new({"a" => ["Apple", "Apple", "Orange"], "b" => [1, nil, 2]})
+    #   lf.group_by("a").len.collect
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌────────┬─────┐
+    #   # │ a      ┆ len │
+    #   # │ ---    ┆ --- │
+    #   # │ str    ┆ u32 │
+    #   # ╞════════╪═════╡
+    #   # │ Apple  ┆ 2   │
+    #   # │ Orange ┆ 1   │
+    #   # └────────┴─────┘
+    #
+    # @example
+    #   lf.group_by("a").len(name: "n").collect
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌────────┬─────┐
+    #   # │ a      ┆ n   │
+    #   # │ ---    ┆ --- │
+    #   # │ str    ┆ u32 │
+    #   # ╞════════╪═════╡
+    #   # │ Apple  ┆ 2   │
+    #   # │ Orange ┆ 1   │
+    #   # └────────┴─────┘
+    def len(name: nil)
+      len_expr = F.len
+      if !name.nil?
+        len_expr = len_expr.alias(name)
+      end
+      agg(len_expr)
+    end
+
+    # Aggregate the first values in the group.
+    #
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default `false`).
+    #   If set to `true`, the first non-null value for each aggregation is returned,
+    #   otherwise `nil` is returned if no non-null value exists.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "c" => [true, true, true, false, false, true],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).first.collect
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌────────┬─────┬──────┬───────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c     │
+    #   # │ ---    ┆ --- ┆ ---  ┆ ---   │
+    #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
+    #   # ╞════════╪═════╪══════╪═══════╡
+    #   # │ Apple  ┆ 1   ┆ 0.5  ┆ true  │
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
+    #   # │ Banana ┆ 4   ┆ 13.0 ┆ false │
+    #   # └────────┴─────┴──────┴───────┘
+    def first(ignore_nulls: false)
+      agg(F.all.first(ignore_nulls: ignore_nulls))
+    end
+
+    # Aggregate the last values in the group.
+    #
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default `false`).
+    #   If set to `true`, the last non-null value for each aggregation is returned,
+    #   otherwise `nil` is returned if no non-null value exists.
+    #
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 14, 13],
+    #       "c" => [true, true, true, false, false, true],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).last.collect
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌────────┬─────┬──────┬───────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c     │
+    #   # │ ---    ┆ --- ┆ ---  ┆ ---   │
+    #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
+    #   # ╞════════╪═════╪══════╪═══════╡
+    #   # │ Apple  ┆ 3   ┆ 10.0 ┆ false │
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
+    #   # │ Banana ┆ 5   ┆ 13.0 ┆ true  │
+    #   # └────────┴─────┴──────┴───────┘
+    def last(ignore_nulls: false)
+      agg(F.all.last(ignore_nulls: ignore_nulls))
+    end
+
+    # Reduce the groups to the maximal value.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "c" => [true, true, true, false, false, true],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).max.collect
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌────────┬─────┬──────┬──────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c    │
+    #   # │ ---    ┆ --- ┆ ---  ┆ ---  │
+    #   # │ str    ┆ i64 ┆ f64  ┆ bool │
+    #   # ╞════════╪═════╪══════╪══════╡
+    #   # │ Apple  ┆ 3   ┆ 10.0 ┆ true │
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ true │
+    #   # │ Banana ┆ 5   ┆ 14.0 ┆ true │
+    #   # └────────┴─────┴──────┴──────┘
+    def max
+      agg(F.all.max)
+    end
+
+    # Reduce the groups to the mean values.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "c" => [true, true, true, false, false, true],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).mean.collect
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌────────┬─────┬──────────┬──────────┐
+    #   # │ d      ┆ a   ┆ b        ┆ c        │
+    #   # │ ---    ┆ --- ┆ ---      ┆ ---      │
+    #   # │ str    ┆ f64 ┆ f64      ┆ f64      │
+    #   # ╞════════╪═════╪══════════╪══════════╡
+    #   # │ Apple  ┆ 2.0 ┆ 4.833333 ┆ 0.666667 │
+    #   # │ Orange ┆ 2.0 ┆ 0.5      ┆ 1.0      │
+    #   # │ Banana ┆ 4.5 ┆ 13.5     ┆ 0.5      │
+    #   # └────────┴─────┴──────────┴──────────┘
+    def mean
+      agg(F.all.mean)
+    end
+
+    # Return the median per group.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "d" => ["Apple", "Banana", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).median.collect
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────┬─────┬──────┐
+    #   # │ d      ┆ a   ┆ b    │
+    #   # │ ---    ┆ --- ┆ ---  │
+    #   # │ str    ┆ f64 ┆ f64  │
+    #   # ╞════════╪═════╪══════╡
+    #   # │ Apple  ┆ 2.0 ┆ 4.0  │
+    #   # │ Banana ┆ 4.0 ┆ 13.0 │
+    #   # └────────┴─────┴──────┘
+    def median
+      agg(F.all.median)
+    end
+
+    # Reduce the groups to the minimal value.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "c" => [true, true, true, false, false, true],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).min.collect
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌────────┬─────┬──────┬───────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c     │
+    #   # │ ---    ┆ --- ┆ ---  ┆ ---   │
+    #   # │ str    ┆ i64 ┆ f64  ┆ bool  │
+    #   # ╞════════╪═════╪══════╪═══════╡
+    #   # │ Apple  ┆ 1   ┆ 0.5  ┆ false │
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ true  │
+    #   # │ Banana ┆ 4   ┆ 13.0 ┆ false │
+    #   # └────────┴─────┴──────┴───────┘
+    def min
+      agg(F.all.min)
+    end
+
+    # Count the unique values per group.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 1, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 0.5, 10, 13, 14],
+    #       "d" => ["Apple", "Banana", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).n_unique.collect
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────┬─────┬─────┐
+    #   # │ d      ┆ a   ┆ b   │
+    #   # │ ---    ┆ --- ┆ --- │
+    #   # │ str    ┆ u32 ┆ u32 │
+    #   # ╞════════╪═════╪═════╡
+    #   # │ Apple  ┆ 2   ┆ 2   │
+    #   # │ Banana ┆ 3   ┆ 3   │
+    #   # └────────┴─────┴─────┘
+    def n_unique
+      agg(F.all.n_unique)
+    end
+
+    # Compute the quantile per group.
+    #
+    # @param quantile [Float]
+    #   Quantile between 0.0 and 1.0.
+    # @param interpolation ['nearest', 'higher', 'lower', 'midpoint', 'linear', 'equiprobable']
+    #   Interpolation method.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).quantile(1).collect
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌────────┬─────┬──────┐
+    #   # │ d      ┆ a   ┆ b    │
+    #   # │ ---    ┆ --- ┆ ---  │
+    #   # │ str    ┆ f64 ┆ f64  │
+    #   # ╞════════╪═════╪══════╡
+    #   # │ Apple  ┆ 3.0 ┆ 10.0 │
+    #   # │ Orange ┆ 2.0 ┆ 0.5  │
+    #   # │ Banana ┆ 5.0 ┆ 14.0 │
+    #   # └────────┴─────┴──────┘
+    def quantile(quantile, interpolation: "nearest")
+      agg(F.all.quantile(quantile, interpolation: interpolation))
+    end
+
+    # Reduce the groups to the sum.
+    #
+    # @return [LazyFrame]
+    #
+    # @example
+    #   ldf = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 2, 2, 3, 4, 5],
+    #       "b" => [0.5, 0.5, 4, 10, 13, 14],
+    #       "c" => [true, true, true, false, false, true],
+    #       "d" => ["Apple", "Orange", "Apple", "Apple", "Banana", "Banana"]
+    #     }
+    #   ).lazy
+    #   ldf.group_by("d", maintain_order: true).sum.collect
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌────────┬─────┬──────┬─────┐
+    #   # │ d      ┆ a   ┆ b    ┆ c   │
+    #   # │ ---    ┆ --- ┆ ---  ┆ --- │
+    #   # │ str    ┆ i64 ┆ f64  ┆ u32 │
+    #   # ╞════════╪═════╪══════╪═════╡
+    #   # │ Apple  ┆ 6   ┆ 14.5 ┆ 2   │
+    #   # │ Orange ┆ 2   ┆ 0.5  ┆ 1   │
+    #   # │ Banana ┆ 9   ┆ 27.0 ┆ 1   │
+    #   # └────────┴─────┴──────┴─────┘
+    def sum
+      agg(F.all.sum)
+    end
   end
 end