polars-df 0.15.0 → 0.26.0

data/lib/polars/list_name_space.rb

@@ -12,6 +12,13 @@ module Polars
 
     # Evaluate whether all boolean values in a list are true.
     #
+    # @param ignore_nulls [Boolean]
+    #   * If set to `true` (default), null values are ignored. If there
+    #     are no non-null values, the output is `true`.
+    #   * If set to `false`, [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic) is used to deal with nulls:
+    #     if the column contains any null values and no `false` values,
+    #     the output is null.
+    #
     # @return [Series]
     #
     # @example
@@ -31,12 +38,19 @@ module Polars
     #   #         true
     #   #         null
     #   # ]
-    def all
+    def all(ignore_nulls: true)
       super
     end
 
     # Evaluate whether any boolean value in a list is true.
     #
+    # @param ignore_nulls [Boolean]
+    #   * If set to `true` (default), null values are ignored. If there
+    #     are no non-null values, the output is `false`.
+    #   * If set to `false`, [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic) is used to deal with nulls:
+    #     if the column contains any null values and no `true` values,
+    #     the output is null.
+    #
     # @return [Series]
     #
     # @example
@@ -56,7 +70,7 @@ module Polars
     #   #         false
     #   #         null
     #   # ]
-    def any
+    def any(ignore_nulls: true)
       super
     end
 
@@ -66,7 +80,7 @@ module Polars
     #
     # @example
     #   s = Polars::Series.new([[1, 2, 3], [5]])
-    #   s.list.lengths
+    #   s.list.len
     #   # =>
     #   # shape: (2,)
     #   # Series: '' [u32]
@@ -74,7 +88,7 @@ module Polars
     #   #         3
     #   #         1
     #   # ]
-    def lengths
+    def len
       super
     end
 
@@ -123,7 +137,7 @@ module Polars
     #   # shape: (2,)
     #   # Series: 'values' [list[i64]]
     #   # [
-    #   #         [2, 1]
+    #   #         [2, 3]
     #   #         [5]
     #   # ]
     def sample(n: nil, fraction: nil, with_replacement: false, shuffle: false, seed: nil)
@@ -202,8 +216,67 @@ module Polars
       super
     end
 
+    # Compute the median value of the arrays in the list.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("values", [[-1, 0, 1], [1, 10]])
+    #   s.list.median
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'values' [f64]
+    #   # [
+    #   #         0.0
+    #   #         5.5
+    #   # ]
+    def median
+      super
+    end
+
+    # Compute the std value of the arrays in the list.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("values", [[-1, 0, 1], [1, 10]])
+    #   s.list.std
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'values' [f64]
+    #   # [
+    #   #         1.0
+    #   #         6.363961
+    #   # ]
+    def std(ddof: 1)
+      super
+    end
+
+    # Compute the var value of the arrays in the list.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("values", [[-1, 0, 1], [1, 10]])
+    #   s.list.var
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'values' [f64]
+    #   # [
+    #   #         1.0
+    #   #         40.5
+    #   # ]
+    def var(ddof: 1)
+      super
+    end
+
     # Sort the arrays in the list.
     #
+    # @param descending [Boolean]
+    #   Sort in descending order.
+    # @param nulls_last [Boolean]
+    #   Place null values last.
+    #
     # @return [Series]
     #
     # @example
@@ -218,7 +291,7 @@ module Polars
     #   # ]
     #
     # @example
-    #   s.list.sort(reverse: true)
+    #   s.list.sort(descending: true)
     #   # =>
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
@@ -226,7 +299,7 @@ module Polars
     #   #         [3, 2, 1]
     #   #         [9, 2, 1]
     #   # ]
-    def sort(reverse: false)
+    def sort(descending: false, nulls_last: false)
       super
     end
 
@@ -250,11 +323,14 @@ module Polars
 
     # Get the unique/distinct values in the list.
     #
+    # @param maintain_order [Boolean]
+    #   Maintain order of data. This requires more work.
+    #
     # @return [Series]
     #
     # @example
     #   s = Polars::Series.new("a", [[1, 1, 2], [2, 3, 3]])
-    #   s.list.unique()
+    #   s.list.unique
     #   # =>
     #   # shape: (2,)
     #   # Series: 'a' [list[i64]]
@@ -262,7 +338,25 @@ module Polars
     #   #         [1, 2]
     #   #         [2, 3]
     #   # ]
-    def unique
+    def unique(maintain_order: false)
+      super
+    end
+
+    # Count the number of unique values in every sub-lists.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 1, 2], [2, 3, 4]])
+    #   s.list.n_unique
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         2
+    #   #         3
+    #   # ]
+    def n_unique
       super
     end
 
@@ -292,7 +386,7 @@ module Polars
     #
     # So index `0` would return the first item of every sublist
     # and index `-1` would return the last item of every sublist
-    # if an index is out of bounds, it will return a `None`.
+    # if an index is out of bounds, it will return a `nil`.
     #
     # @param index [Integer]
     #   Index to return per sublist
@@ -318,6 +412,63 @@ module Polars
       super
     end
 
+    # Take sublists by multiple indices.
+    #
+    # The indices may be defined in a single column, or by sublists in another
+    # column of dtype `List`.
+    #
+    # @param indices [Object]
+    #   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
+    #   Note that defaulting to raising an error is much cheaper
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[3, 2, 1], [], [1, 2]])
+    #   s.list.gather([0, 2], null_on_oob: true)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [list[i64]]
+    #   # [
+    #   #         [3, 1]
+    #   #         [null, null]
+    #   #         [1, null]
+    #   # ]
+    def gather(
+      indices,
+      null_on_oob: false
+    )
+      super
+    end
+
+    # Take every n-th value start from offset in sublists.
+    #
+    # @param n [Integer]
+    #   Gather every n-th element.
+    # @param offset [Integer]
+    #   Starting index.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2, 3], [], [6, 7, 8, 9]])
+    #   s.list.gather_every(2, 1)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [list[i64]]
+    #   # [
+    #   #         [2]
+    #   #         []
+    #   #         [7, 9]
+    #   # ]
+    def gather_every(n, offset = 0)
+      super
+    end
+
     # Get the value by index in the sublists.
     #
     # @return [Series]
@@ -331,6 +482,11 @@ module Polars
     #
     # @param separator [String]
     #   string to separate the items with
+    # @param ignore_nulls [Boolean]
+    #   Ignore null values (default).
+    #
+    #   If set to `false`, null values will be propagated.
+    #   If the sub-list contains any null values, the output is `nil`.
     #
     # @return [Series]
     #
@@ -344,7 +500,7 @@ module Polars
     #   #         "foo-bar"
     #   #         "hello-world"
     #   # ]
-    def join(separator)
+    def join(separator, ignore_nulls: true)
       super
     end
 
@@ -386,10 +542,33 @@ module Polars
       super
     end
 
+    # Get the single value of the sublists.
+    #
+    # This errors if the sublist length is not exactly one.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1], [4], [6]])
+    #   s.list.item
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         1
+    #   #         4
+    #   #         6
+    #   # ]
+    def item
+      super
+    end
+
     # Check if sublists contain the given item.
     #
     # @param item [Object]
     #   Item that will be checked for membership.
+    # @param nulls_equal [Boolean]
+    #   If true, treat null as a distinct value. Null values will not propagate.
     #
     # @return [Series]
     #
@@ -404,7 +583,7 @@ module Polars
     #   #         false
     #   #         true
     #   # ]
-    def contains(item)
+    def contains(item, nulls_equal: true)
       super
     end
 
@@ -469,7 +648,7 @@ module Polars
 
     # Shift values by the given period.
     #
-    # @param periods [Integer]
+    # @param n [Integer]
     #   Number of places to shift (may be negative).
     #
     # @return [Series]
@@ -484,7 +663,7 @@ module Polars
     #   #         [null, 1, … 3]
     #   #         [null, 10, 2]
     #   # ]
-    def shift(periods = 1)
+    def shift(n = 1)
       super
     end
 
@@ -554,31 +733,124 @@ module Polars
       super
     end
 
+    # Returns a column with a separate row for every list element.
+    #
+    # @param empty_as_null [Boolean]
+    #   Explode an empty list into a `null`.
+    # @param keep_nulls [Boolean]
+    #   Explode a `null` list into a `null`.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 2, 3], [4, 5, 6]])
+    #   s.list.explode
+    #   # =>
+    #   # shape: (6,)
+    #   # Series: 'a' [i64]
+    #   # [
+    #   #         1
+    #   #         2
+    #   #         3
+    #   #         4
+    #   #         5
+    #   #         6
+    #   # ]
+    def explode(empty_as_null: true, keep_nulls: true)
+      super
+    end
+
+    # Count how often the value produced by `element` occurs.
+    #
+    # @param element [Object]
+    #   An expression that produces a single value
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[0], [1], [1, 2, 3, 2], [1, 2, 1], [4, 4]])
+    #   s.list.count_matches(1)
+    #   # =>
+    #   # shape: (5,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         0
+    #   #         1
+    #   #         1
+    #   #         2
+    #   #         0
+    #   # ]
+    def count_matches(element)
+      super
+    end
+
+    # Convert a List column into an Array column with the same inner data type.
+    #
+    # @param width [Integer]
+    #   Width of the resulting Array column.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new([[1, 2], [3, 4]], dtype: Polars::List.new(Polars::Int8))
+    #   s.list.to_array(2)
+    #   # =>
+    #   # shape: (2,)
+    #   # Series: '' [array[i8, 2]]
+    #   # [
+    #   #         [1, 2]
+    #   #         [3, 4]
+    #   # ]
+    def to_array(width)
+      super
+    end
+
     # Convert the series of type `List` to a series of type `Struct`.
     #
     # @param n_field_strategy ["first_non_null", "max_width"]
     #   Strategy to determine the number of fields of the struct.
-    # @param name_generator [Object]
-    #   A custom function that can be used to generate the field names.
-    #   Default field names are `field_0, field_1 .. field_n`
+    # @param fields [Object]
+    #   If the name and number of the desired fields is known in advance
+    #   a list of field names can be given, which will be assigned by index.
+    #   Otherwise, to dynamically assign field names, a custom function can be
+    #   used; if neither are set, fields will be `field_0, field_1 .. field_n`.
     #
     # @return [Series]
     #
-    # @example
-    #   df = Polars::DataFrame.new({"a" => [[1, 2, 3], [1, 2]]})
-    #   df.select([Polars.col("a").list.to_struct])
+    # @example Convert list to struct with field name assignment by index from a list of names:
+    #   s1 = Polars::Series.new("n", [[0, 1, 2], [0, 1]])
+    #   s2 = s1.list.to_struct
+    #   s2.struct.fields
+    #   # => ["field_0", "field_1", "field_2"]
+    #
+    # @example Convert list to struct with field name assignment by function/index:
+    #   s3 = s1.list.to_struct(fields: ->(idx) { "n%02d" % idx })
+    #   s3.struct.fields
+    #   # => ["n00", "n01", "n02"]
+    #
+    # @example Convert list to struct with field name assignment by index from a list of names:
+    #   s1.list.to_struct(fields: ["one", "two", "three"]).struct.unnest
     #   # =>
-    #   # shape: (2, 1)
-    #   # ┌────────────┐
-    #   # │ a          │
-    #   # │ ---        │
-    #   # │ struct[3]  │
-    #   # ╞════════════╡
-    #   # │ {1,2,3}    │
-    #   # │ {1,2,null} │
-    #   # └────────────┘
-    def to_struct(n_field_strategy: "first_non_null", name_generator: nil)
-      super
+    #   # shape: (2, 3)
+    #   # ┌─────┬─────┬───────┐
+    #   # │ one ┆ two ┆ three │
+    #   # │ --- ┆ --- ┆ ---   │
+    #   # │ i64 ┆ i64 ┆ i64   │
+    #   # ╞═════╪═════╪═══════╡
+    #   # │ 0   ┆ 1   ┆ 2     │
+    #   # │ 0   ┆ 1   ┆ null  │
+    #   # └─────┴─────┴───────┘
+    def to_struct(n_field_strategy: "first_non_null", fields: nil)
+      if fields.is_a?(::Array)
+        s = Utils.wrap_s(_s)
+        return (
+          s.to_frame
+          .select_seq(F.col(s.name).list.to_struct(fields: fields))
+          .to_series
+        )
+      end
+
+      Utils.wrap_s(_s.list_to_struct(n_field_strategy, fields))
     end
 
     # Run any polars expression against the lists' elements.
@@ -586,32 +858,174 @@ module Polars
     # @param expr [Expr]
     #   Expression to run. Note that you can select an element with `Polars.first`, or
     #   `Polars.col`
-    # @param parallel [Boolean]
-    #   Run all expression parallel. Don't activate this blindly.
-    #   Parallelism is worth it if there is enough work to do per thread.
     #
-    #   This likely should not be use in the group by context, because we already
-    #   parallel execution per group
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 4], [8, 5], [3, 2]])
+    #   s.list.eval(Polars.element.rank)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [list[f64]]
+    #   # [
+    #   #         [1.0, 2.0]
+    #   #         [2.0, 1.0]
+    #   #         [2.0, 1.0]
+    #   # ]
+    def eval(expr)
+      s = Utils.wrap_s(_s)
+      s.to_frame.select(F.col(s.name).list.eval(expr)).to_series
+    end
+
+    # Run any polars aggregation expression against the list' elements.
+    #
+    # @param expr [Expr]
+    #   Expression to run. Note that you can select an element with `Polars.element`.
     #
     # @return [Series]
     #
     # @example
-    #   df = Polars::DataFrame.new({"a" => [1, 8, 3], "b" => [4, 5, 2]})
-    #   df.with_column(
-    #     Polars.concat_list(["a", "b"]).list.eval(Polars.element.rank).alias("rank")
-    #   )
+    #   s = Polars::Series.new("a", [[1, nil], [42, 13], [nil, nil]])
+    #   s.list.agg(Polars.element.null_count)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [u32]
+    #   # [
+    #   #         1
+    #   #         0
+    #   #         2
+    #   # ]
+    #
+    # @example
+    #   s.list.agg(Polars.element.drop_nulls)
     #   # =>
-    #   # shape: (3, 3)
-    #   # ┌─────┬─────┬────────────┐
-    #   # │ a   ┆ b   ┆ rank       │
-    #   # │ --- ┆ --- ┆ ---        │
-    #   # │ i64 ┆ i64 ┆ list[f64]  │
-    #   # ╞═════╪═════╪════════════╡
-    #   # │ 1   ┆ 4   ┆ [1.0, 2.0] │
-    #   # │ 8   ┆ 5   ┆ [2.0, 1.0] │
-    #   # │ 3   ┆ 2   ┆ [2.0, 1.0] │
-    #   # └─────┴─────┴────────────┘
-    def eval(expr, parallel: false)
+    #   # shape: (3,)
+    #   # Series: 'a' [list[i64]]
+    #   # [
+    #   #         [1]
+    #   #         [42, 13]
+    #   #         []
+    #   # ]
+    def agg(expr)
+      super
+    end
+
+    # Filter elements in each list by a boolean expression, returning a new Series of lists.
+    #
+    # @param predicate [Object]
+    #   A boolean expression evaluated on each list element.
+    #   Use `Polars.element` to refer to the current element.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   s = Polars::Series.new("a", [[1, 4], [8, 5], [3, 2]])
+    #   s.list.filter(Polars.element % 2 == 0)
+    #   # =>
+    #   # shape: (3,)
+    #   # Series: 'a' [list[i64]]
+    #   # [
+    #   #         [4]
+    #   #         [8]
+    #   #         [2]
+    #   # ]
+    def filter(predicate)
+      super
+    end
+
+    # Compute the SET UNION between the elements in this list and the elements of `other`.
+    #
+    # @param other [Object]
+    #   Right hand side of the set operation.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   a = Polars::Series.new([[1, 2, 3], [], [nil, 3], [5, 6, 7]])
+    #   b = Polars::Series.new([[2, 3, 4], [3], [3, 4, nil], [6, 8]])
+    #   a.list.set_union(b)
+    #   # =>
+    #   # shape: (4,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [1, 2, … 4]
+    #   #         [3]
+    #   #         [null, 3, 4]
+    #   #         [5, 6, … 8]
+    #   # ]
+    def set_union(other)
+      super
+    end
+
+    # Compute the SET DIFFERENCE between the elements in this list and the elements of `other`.
+    #
+    # @param other [Object]
+    #   Right hand side of the set operation.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   a = Polars::Series.new([[1, 2, 3], [], [nil, 3], [5, 6, 7]])
+    #   b = Polars::Series.new([[2, 3, 4], [3], [3, 4, nil], [6, 8]])
+    #   a.list.set_difference(b)
+    #   # =>
+    #   # shape: (4,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [1]
+    #   #         []
+    #   #         []
+    #   #         [5, 7]
+    #   # ]
+    def set_difference(other)
+      super
+    end
+
+    # Compute the SET INTERSECTION between the elements in this list and the elements of `other`.
+    #
+    # @param other [Object]
+    #   Right hand side of the set operation.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   a = Polars::Series.new([[1, 2, 3], [], [nil, 3], [5, 6, 7]])
+    #   b = Polars::Series.new([[2, 3, 4], [3], [3, 4, nil], [6, 8]])
+    #   a.list.set_intersection(b)
+    #   # =>
+    #   # shape: (4,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [2, 3]
+    #   #         []
+    #   #         [null, 3]
+    #   #         [6]
+    #   # ]
+    def set_intersection(other)
+      super
+    end
+
+    # Compute the SET SYMMETRIC DIFFERENCE between the elements in this list and the elements of `other`.
+    #
+    # @param other [Object]
+    #   Right hand side of the set operation.
+    #
+    # @return [Series]
+    #
+    # @example
+    #   a = Polars::Series.new([[1, 2, 3], [], [nil, 3], [5, 6, 7]])
+    #   b = Polars::Series.new([[2, 3, 4], [3], [3, 4, nil], [6, 8]])
+    #   a.list.set_symmetric_difference(b)
+    #   # =>
+    #   # shape: (4,)
+    #   # Series: '' [list[i64]]
+    #   # [
+    #   #         [1, 4]
+    #   #         [3]
+    #   #         [4]
+    #   #         [5, 7, 8]
+    #   # ]
+    def set_symmetric_difference(other)
       super
     end
   end