polars-df 0.23.0 → 0.24.0

data/lib/polars/lazy_group_by.rb

@@ -175,7 +175,332 @@ 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.
+    #
+    # @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
+      agg(F.all.first)
+    end
+
+    # Aggregate the last values in the group.
+    #
+    # @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
+      agg(F.all.last)
+    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

data/lib/polars/list_expr.rb

@@ -69,7 +69,7 @@ module Polars
     #
     # @example
     #   df = Polars::DataFrame.new({"foo" => [1, 2], "bar" => [["a", "b"], ["c"]]})
-    #   df.select(Polars.col("bar").list.lengths)
+    #   df.select(Polars.col("bar").list.len)
     #   # =>
     #   # shape: (2, 1)
     #   # ┌─────┐
@@ -83,7 +83,6 @@ module Polars
     def len
       Utils.wrap_expr(_rbexpr.list_len)
     end
-    alias_method :lengths, :len
 
     # Drop all null values in the list.
     #
@@ -318,7 +317,7 @@ module Polars
 
     # Sort the arrays in the list.
     #
-    # @param reverse [Boolean]
+    # @param descending [Boolean]
     #   Sort in descending order.
     # @param nulls_last [Boolean]
     #   Place null values last.
@@ -342,8 +341,8 @@ module Polars
     #   # │ [1, 2, 3] │
     #   # │ [1, 2, 9] │
     #   # └───────────┘
-    def sort(reverse: false, nulls_last: false)
-      Utils.wrap_expr(_rbexpr.list_sort(reverse, nulls_last))
+    def sort(descending: false, nulls_last: false)
+      Utils.wrap_expr(_rbexpr.list_sort(descending, nulls_last))
     end
 
     # Reverse the arrays in the list.
@@ -477,7 +476,7 @@ module Polars
     #
     # @example
     #   df = Polars::DataFrame.new({"foo" => [[3, 2, 1], [], [1, 2]]})
-    #   df.select(Polars.col("foo").list.get(0))
+    #   df.select(Polars.col("foo").list.get(0, null_on_oob: true))
     #   # =>
     #   # shape: (3, 1)
     #   # ┌──────┐
@@ -489,7 +488,7 @@ module Polars
     #   # │ null │
     #   # │ 1    │
     #   # └──────┘
-    def get(index, null_on_oob: true)
+    def get(index, null_on_oob: false)
       index = Utils.parse_into_expression(index)
       Utils.wrap_expr(_rbexpr.list_get(index, null_on_oob))
     end
@@ -510,8 +509,8 @@ module Polars
     #   Indices to return per sublist
     # @param null_on_oob [Boolean]
     #   Behavior if an index is out of bounds:
-    #   True -> set as null
-    #   False -> raise an error
+    #   true -> set as null
+    #   false -> raise an error
     #   Note that defaulting to raising an error is much cheaper
     #
     # @return [Expr]
@@ -534,7 +533,6 @@ module Polars
       indices = Utils.parse_into_expression(indices)
       Utils.wrap_expr(_rbexpr.list_gather(indices, null_on_oob))
     end
-    alias_method :take, :gather
 
     # Take every n-th value start from offset in sublists.
     #
@@ -597,7 +595,7 @@ module Polars
     #   # │ 1    │
     #   # └──────┘
     def first
-      get(0)
+      get(0, null_on_oob: true)
     end
 
     # Get the last value of the sublists.
@@ -619,7 +617,49 @@ module Polars
     #   # │ 2    │
     #   # └──────┘
     def last
-      get(-1)
+      get(-1, null_on_oob: true)
+    end
+
+    # Get the single value of the sublists.
+    #
+    # This errors if the sublist length is not exactly one.
+    #
+    # @param allow_empty [Boolean]
+    #   Allow having no values to return `null`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [[3], [1], [2]]})
+    #   df.with_columns(item: Polars.col("a").list.item)
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌───────────┬──────┐
+    #   # │ a         ┆ item │
+    #   # │ ---       ┆ ---  │
+    #   # │ list[i64] ┆ i64  │
+    #   # ╞═══════════╪══════╡
+    #   # │ [3]       ┆ 3    │
+    #   # │ [1]       ┆ 1    │
+    #   # │ [2]       ┆ 2    │
+    #   # └───────────┴──────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [[], [1], [2]]})
+    #   df.select(Polars.col("a").list.item(allow_empty: true))
+    #   # =>
+    #   # shape: (3, 1)
+    #   # ┌──────┐
+    #   # │ a    │
+    #   # │ ---  │
+    #   # │ i64  │
+    #   # ╞══════╡
+    #   # │ null │
+    #   # │ 1    │
+    #   # │ 2    │
+    #   # └──────┘
+    def item(allow_empty: false)
+      agg(F.element.item(allow_empty: allow_empty))
     end
 
     # Check if sublists contain the given item.
@@ -876,7 +916,7 @@ module Polars
     #
     # @example
     #   df = Polars::DataFrame.new({"listcol" => [[0], [1], [1, 2, 3, 2], [1, 2, 1], [4, 4]]})
-    #   df.select(Polars.col("listcol").list.count_match(2).alias("number_of_twos"))
+    #   df.select(Polars.col("listcol").list.count_matches(2).alias("number_of_twos"))
     #   # =>
     #   # shape: (5, 1)
     #   # ┌────────────────┐
@@ -893,7 +933,6 @@ module Polars
     def count_matches(element)
       Utils.wrap_expr(_rbexpr.list_count_matches(Utils.parse_into_expression(element)))
     end
-    alias_method :count_match, :count_matches
 
     # Convert a List column into an Array column with the same inner data type.
     #
@@ -957,7 +996,7 @@ module Polars
     #   # │ [0, 1]    ┆ {0,1}     │
     #   # │ [0, 1, 2] ┆ {0,1}     │
     #   # └───────────┴───────────┘
-    def to_struct(n_field_strategy: "first_non_null", fields: nil, upper_bound: nil)
+    def to_struct(n_field_strategy: nil, fields: nil, upper_bound: nil)
       if !fields.is_a?(::Array)
         if fields.nil?
           fields = upper_bound.times.map { |i| "field_#{i}" }
@@ -979,7 +1018,7 @@ module Polars
     #
     # @example
     #   df = Polars::DataFrame.new({"a" => [1, 8, 3], "b" => [4, 5, 2]})
-    #   df.with_column(
+    #   df.with_columns(
     #     Polars.concat_list(["a", "b"]).list.eval(Polars.element.rank).alias("rank")
     #   )
     #   # =>
@@ -997,6 +1036,45 @@ module Polars
       Utils.wrap_expr(_rbexpr.list_eval(expr._rbexpr))
     end
 
+    # Run any polars aggregation expression against the lists' elements.
+    #
+    # @param expr [Expr]
+    #   Expression to run. Note that you can select an element with `Polars.element`.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [[1, nil], [42, 13], [nil, nil]]})
+    #   df.with_columns(null_count: Polars.col("a").list.agg(Polars.element.null_count))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌──────────────┬────────────┐
+    #   # │ a            ┆ null_count │
+    #   # │ ---          ┆ ---        │
+    #   # │ list[i64]    ┆ u32        │
+    #   # ╞══════════════╪════════════╡
+    #   # │ [1, null]    ┆ 1          │
+    #   # │ [42, 13]     ┆ 0          │
+    #   # │ [null, null] ┆ 2          │
+    #   # └──────────────┴────────────┘
+    #
+    # @example
+    #   df.with_columns(no_nulls: Polars.col("a").list.agg(Polars.element.drop_nulls))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌──────────────┬───────────┐
+    #   # │ a            ┆ no_nulls  │
+    #   # │ ---          ┆ ---       │
+    #   # │ list[i64]    ┆ list[i64] │
+    #   # ╞══════════════╪═══════════╡
+    #   # │ [1, null]    ┆ [1]       │
+    #   # │ [42, 13]     ┆ [42, 13]  │
+    #   # │ [null, null] ┆ []        │
+    #   # └──────────────┴───────────┘
+    def agg(expr)
+      Utils.wrap_expr(_rbexpr.list_agg(expr._rbexpr))
+    end
+
     # Filter elements in each list by a boolean expression.
     #
     # @param predicate [Object]