polars-df 0.8.0-x86_64-linux → 0.9.0-x86_64-linux

data/lib/polars/list_expr.rb

@@ -9,6 +9,60 @@ module Polars
       self._rbexpr = expr._rbexpr
     end
 
+    # Evaluate whether all boolean values in a list are true.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[true, true], [false, true], [false, false], [nil], [], nil]}
+    #   )
+    #   df.with_columns(all: Polars.col("a").list.all)
+    #   # =>
+    #   # shape: (6, 2)
+    #   # ┌────────────────┬───────┐
+    #   # │ a              ┆ all   │
+    #   # │ ---            ┆ ---   │
+    #   # │ list[bool]     ┆ bool  │
+    #   # ╞════════════════╪═══════╡
+    #   # │ [true, true]   ┆ true  │
+    #   # │ [false, true]  ┆ false │
+    #   # │ [false, false] ┆ false │
+    #   # │ [null]         ┆ true  │
+    #   # │ []             ┆ true  │
+    #   # │ null           ┆ null  │
+    #   # └────────────────┴───────┘
+    def all
+      Utils.wrap_expr(_rbexpr.list_all)
+    end
+
+    # Evaluate whether any boolean value in a list is true.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[true, true], [false, true], [false, false], [nil], [], nil]}
+    #   )
+    #   df.with_columns(any: Polars.col("a").list.any)
+    #   # =>
+    #   # shape: (6, 2)
+    #   # ┌────────────────┬───────┐
+    #   # │ a              ┆ any   │
+    #   # │ ---            ┆ ---   │
+    #   # │ list[bool]     ┆ bool  │
+    #   # ╞════════════════╪═══════╡
+    #   # │ [true, true]   ┆ true  │
+    #   # │ [false, true]  ┆ true  │
+    #   # │ [false, false] ┆ false │
+    #   # │ [null]         ┆ false │
+    #   # │ []             ┆ false │
+    #   # │ null           ┆ null  │
+    #   # └────────────────┴───────┘
+    def any
+      Utils.wrap_expr(_rbexpr.list_any)
+    end
+
     # Get the length of the arrays as `:u32`.
     #
     # @return [Expr]
@@ -31,6 +85,80 @@ module Polars
     end
     alias_method :len, :lengths
 
+    # Drop all null values in the list.
+    #
+    # The original order of the remaining elements is preserved.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"values" => [[nil, 1, nil, 2], [nil], [3, 4]]})
+    #   df.with_columns(drop_nulls: Polars.col("values").list.drop_nulls)
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌────────────────┬────────────┐
+    #   # │ values         ┆ drop_nulls │
+    #   # │ ---            ┆ ---        │
+    #   # │ list[i64]      ┆ list[i64]  │
+    #   # ╞════════════════╪════════════╡
+    #   # │ [null, 1, … 2] ┆ [1, 2]     │
+    #   # │ [null]         ┆ []         │
+    #   # │ [3, 4]         ┆ [3, 4]     │
+    #   # └────────────────┴────────────┘
+    def drop_nulls
+      Utils.wrap_expr(_rbexpr.list_drop_nulls)
+    end
+
+    # Sample from this list.
+    #
+    # @param n [Integer]
+    #   Number of items to return. Cannot be used with `fraction`. Defaults to 1 if
+    #   `fraction` is nil.
+    # @param fraction [Float]
+    #   Fraction of items to return. Cannot be used with `n`.
+    # @param with_replacement [Boolean]
+    #   Allow values to be sampled more than once.
+    # @param shuffle [Boolean]
+    #   Shuffle the order of sampled data points.
+    # @param seed [Integer]
+    #   Seed for the random number generator. If set to nil (default), a
+    #   random seed is generated for each sample operation.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"values" => [[1, 2, 3], [4, 5]], "n" => [2, 1]})
+    #   df.with_columns(sample: Polars.col("values").list.sample(n: Polars.col("n"), seed: 1))
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌───────────┬─────┬───────────┐
+    #   # │ values    ┆ n   ┆ sample    │
+    #   # │ ---       ┆ --- ┆ ---       │
+    #   # │ list[i64] ┆ i64 ┆ list[i64] │
+    #   # ╞═══════════╪═════╪═══════════╡
+    #   # │ [1, 2, 3] ┆ 2   ┆ [2, 1]    │
+    #   # │ [4, 5]    ┆ 1   ┆ [5]       │
+    #   # └───────────┴─────┴───────────┘
+    def sample(n: nil, fraction: nil, with_replacement: false, shuffle: false, seed: nil)
+      if !n.nil? && !fraction.nil?
+        msg = "cannot specify both `n` and `fraction`"
+        raise ArgumentError, msg
+      end
+
+      if !fraction.nil?
+        fraction = Utils.parse_as_expression(fraction)
+        return Utils.wrap_expr(
+          _rbexpr.list_sample_fraction(
+            fraction, with_replacement, shuffle, seed
+          )
+        )
+      end
+
+      n = 1 if n.nil?
+      n = Utils.parse_as_expression(n)
+      Utils.wrap_expr(_rbexpr.list_sample_n(n, with_replacement, shuffle, seed))
+    end
+
     # Sum all the lists in the array.
     #
     # @return [Expr]
@@ -280,13 +408,29 @@ module Polars
     #   Note that defaulting to raising an error is much cheaper
     #
     # @return [Expr]
-    def take(index, null_on_oob: false)
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"a" => [[3, 2, 1], [], [1, 2, 3, 4, 5]]})
+    #   df.with_columns(gather: Polars.col("a").list.gather([0, 4], null_on_oob: true))
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────────────┬──────────────┐
+    #   # │ a           ┆ gather       │
+    #   # │ ---         ┆ ---          │
+    #   # │ list[i64]   ┆ list[i64]    │
+    #   # ╞═════════════╪══════════════╡
+    #   # │ [3, 2, 1]   ┆ [3, null]    │
+    #   # │ []          ┆ [null, null] │
+    #   # │ [1, 2, … 5] ┆ [1, 5]       │
+    #   # └─────────────┴──────────────┘
+    def gather(index, null_on_oob: false)
       if index.is_a?(::Array)
         index = Series.new(index)
       end
       index = Utils.expr_to_lit_or_expr(index, str_to_lit: false)._rbexpr
-      Utils.wrap_expr(_rbexpr.list_take(index, null_on_oob))
+      Utils.wrap_expr(_rbexpr.list_gather(index, null_on_oob))
     end
+    alias_method :take, :gather
 
     # Get the first value of the sublists.
     #
@@ -363,6 +507,8 @@ module Polars
     #
     # @param separator [String]
     #   string to separate the items with
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default).
     #
     # @return [Expr]
     #
@@ -379,9 +525,9 @@ module Polars
     #   # │ a b c │
     #   # │ x y   │
     #   # └───────┘
-    def join(separator)
+    def join(separator, ignore_nulls: true)
       separator = Utils.parse_as_expression(separator, str_as_lit: true)
-      Utils.wrap_expr(_rbexpr.list_join(separator))
+      Utils.wrap_expr(_rbexpr.list_join(separator, ignore_nulls))
     end
 
     # Retrieve the index of the minimal value in every sublist.
@@ -544,8 +690,8 @@ module Polars
     #   #         [2, 1]
     #   # ]
     def tail(n = 5)
-      offset = -Utils.expr_to_lit_or_expr(n, str_to_lit: false)
-      slice(offset, n)
+      n = Utils.parse_as_expression(n)
+      Utils.wrap_expr(_rbexpr.list_tail(n))
     end
 
     # Count how often the value produced by ``element`` occurs.

data/lib/polars/list_name_space.rb

@@ -10,6 +10,56 @@ module Polars
       self._s = series._s
     end
 
+    # Evaluate whether all boolean values in a list are true.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     [[true, true], [false, true], [false, false], [nil], [], nil],
+    #     dtype: Polars::List.new(Polars::Boolean)
+    #   )
+    #   s.list.all
+    #   # =>
+    #   # shape: (6,)
+    #   # Series: '' [bool]
+    #   # [
+    #   #         true
+    #   #         false
+    #   #         false
+    #   #         true
+    #   #         true
+    #   #         null
+    #   # ]
+    def all
+      super
+    end
+
+    # Evaluate whether any boolean value in a list is true.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new(
+    #     [[true, true], [false, true], [false, false], [nil], [], nil],
+    #     dtype: Polars::List.new(Polars::Boolean)
+    #   )
+    #   s.list.any
+    #   # =>
+    #   # shape: (6,)
+    #   # Series: '' [bool]
+    #   # [
+    #   #         true
+    #   #         true
+    #   #         false
+    #   #         false
+    #   #         false
+    #   #         null
+    #   # ]
+    def any
+      super
+    end
+
     # Get the length of the arrays as UInt32.
     #
     # @return [Series]
@@ -28,6 +78,58 @@ module Polars
       super
     end
 
+    # Drop all null values in the list.
+    #
+    # The original order of the remaining elements is preserved.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("values", [[nil, 1, nil, 2], [nil], [3, 4]])
+    #   s.list.drop_nulls
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'values' [list[i64]]
+    #   # [
+    #   #         [1, 2]
+    #   #         []
+    #   #         [3, 4]
+    #   # ]
+    def drop_nulls
+      super
+    end
+
+    # Sample from this list.
+    #
+    # @param n [Integer]
+    #   Number of items to return. Cannot be used with `fraction`. Defaults to 1 if
+    #   `fraction` is nil.
+    # @param fraction [Float]
+    #   Fraction of items to return. Cannot be used with `n`.
+    # @param with_replacement [Boolean]
+    #   Allow values to be sampled more than once.
+    # @param shuffle [Boolean]
+    #   Shuffle the order of sampled data points.
+    # @param seed [Integer]
+    #   Seed for the random number generator. If set to nil (default), a
+    #   random seed is generated for each sample operation.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("values", [[1, 2, 3], [4, 5]])
+    #   s.list.sample(n: Polars::Series.new("n", [2, 1]), seed: 1)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'values' [list[i64]]
+    #   # [
+    #   #         [2, 1]
+    #   #         [5]
+    #   # ]
+    def sample(n: nil, fraction: nil, with_replacement: false, shuffle: false, seed: nil)
+      super
+    end
+
     # Sum all the arrays in the list.
     #
     # @return [Series]

data/lib/polars/meta_expr.rb

@@ -23,9 +23,114 @@ module Polars
       !(self == other)
     end
 
+    # Indicate if this expression is the same as another expression.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   foo_bar = Polars.col("foo").alias("bar")
+    #   foo = Polars.col("foo")
+    #   foo_bar.meta.eq(foo)
+    #   # => false
+    #   foo_bar2 = Polars.col("foo").alias("bar")
+    #   foo_bar.meta.eq(foo_bar2)
+    #   # => true
+    def eq(other)
+      _rbexpr.meta_eq(other._rbexpr)
+    end
+
+    # Indicate if this expression is NOT the same as another expression.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   foo_bar = Polars.col("foo").alias("bar")
+    #   foo = Polars.col("foo")
+    #   foo_bar.meta.ne(foo)
+    #   # => true
+    #   foo_bar2 = Polars.col("foo").alias("bar")
+    #   foo_bar.meta.ne(foo_bar2)
+    #   # => false
+    def ne(other)
+      !eq(other)
+    end
+
+    # Indicate if this expression expands into multiple expressions.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   e = Polars.col(["a", "b"]).alias("bar")
+    #   e.meta.has_multiple_outputs
+    #   # => true
+    def has_multiple_outputs
+      _rbexpr.meta_has_multiple_outputs
+    end
+
+    # Indicate if this expression is a basic (non-regex) unaliased column.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   e = Polars.col("foo")
+    #   e.meta.is_column
+    #   # => true
+    #   e = Polars.col("foo") * Polars.col("bar")
+    #   e.meta.is_column
+    #   # => false
+    #   e = Polars.col("^col.*\d+$")
+    #   e.meta.is_column
+    #   # => false
+    def is_column
+      _rbexpr.meta_is_column
+    end
+
+    # Indicate if this expression expands to columns that match a regex pattern.
+    #
+    # @return [Boolean]
+    #
+    # @example
+    #   e = Polars.col("^.*$").alias("bar")
+    #   e.meta.is_regex_projection
+    #   # => true
+    def is_regex_projection
+      _rbexpr.meta_is_regex_projection
+    end
+
+    # Get the column name that this expression would produce.
+    #
+    # @return [String]
+    #
+    # @example
+    #   e = Polars.col("foo") * Polars.col("bar")
+    #   e.meta.output_name
+    #   # => "foo"
+    #   e_filter = Polars.col("foo").filter(Polars.col("bar") == 13)
+    #   e_filter.meta.output_name
+    #   # => "foo"
+    #   e_sum_over = Polars.sum("foo").over("groups")
+    #   e_sum_over.meta.output_name
+    #   # => "foo"
+    #   e_sum_slice = Polars.sum("foo").slice(Polars.len - 10, Polars.col("bar"))
+    #   e_sum_slice.meta.output_name
+    #   # => "foo"
+    #   Polars.len.meta.output_name
+    #   # => "len"
+    def output_name
+      _rbexpr.meta_output_name
+    end
+
     # Pop the latest expression and return the input(s) of the popped expression.
     #
     # @return [Array]
+    #
+    # @example
+    #   e = Polars.col("foo").alias("bar")
+    #   first = e.meta.pop[0]
+    #   _ = first.meta == Polars.col("foo")
+    #   # => true
+    #   _ = first.meta == Polars.col("bar")
+    #   # => false
     def pop
       _rbexpr.meta_pop.map { |e| Utils.wrap_expr(e) }
     end
@@ -33,22 +138,85 @@ module Polars
     # Get a list with the root column name.
     #
     # @return [Array]
+    #
+    # @example
+    #   e = Polars.col("foo") * Polars.col("bar")
+    #   e.meta.root_names
+    #   # => ["foo", "bar"]
+    #   e_filter = Polars.col("foo").filter(Polars.col("bar") == 13)
+    #   e_filter.meta.root_names
+    #   # => ["foo", "bar"]
+    #   e_sum_over = Polars.sum("foo").over("groups")
+    #   e_sum_over.meta.root_names
+    #   # => ["foo", "groups"]
+    #   e_sum_slice = Polars.sum("foo").slice(Polars.len - 10, Polars.col("bar"))
+    #   e_sum_slice.meta.root_names
+    #   # => ["foo", "bar"]
     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]
+    #
+    # @example
+    #   e = Polars.col("foo").alias("bar")
+    #   _ = e.meta.undo_aliases.meta == Polars.col("foo")
+    #   # => true
+    #   e = Polars.col("foo").sum.over("bar")
+    #   _ = e.name.keep.meta.undo_aliases.meta == e
+    #   # => true
     def undo_aliases
       Utils.wrap_expr(_rbexpr.meta_undo_aliases)
     end
+
+    # Turn this expression in a selector.
+    #
+    # @return [Expr]
+    def _as_selector
+      Utils.wrap_expr(_rbexpr._meta_as_selector)
+    end
+
+    # Add selectors.
+    #
+    # @return [Expr]
+    def _selector_add(other)
+      Utils.wrap_expr(_rbexpr._meta_selector_add(other._rbexpr))
+    end
+
+    # Subtract selectors.
+    #
+    # @return [Expr]
+    def _selector_sub(other)
+      Utils.wrap_expr(_rbexpr._meta_selector_sub(other._rbexpr))
+    end
+
+    # & selectors.
+    #
+    # @return [Expr]
+    def _selector_and(other)
+      Utils.wrap_expr(_rbexpr._meta_selector_and(other._rbexpr))
+    end
+
+    # Format the expression as a tree.
+    #
+    # @param return_as_string [Boolean]
+    #   If true, return as string rather than printing to stdout.
+    #
+    # @return [String]
+    #
+    # @example
+    #   e = (Polars.col("foo") * Polars.col("bar")).sum.over(Polars.col("ham")) / 2
+    #   e.meta.tree_format(return_as_string: true)
+    def tree_format(return_as_string: false)
+      s = _rbexpr.meta_tree_format
+      if return_as_string
+        s
+      else
+        puts s
+        nil
+      end
+    end
   end
 end