polars-df 0.21.0 → 0.21.1

data/lib/polars/array_expr.rb

@@ -9,6 +9,181 @@ module Polars
       self._rbexpr = expr._rbexpr
     end
 
+    # Return the number of elements in each array.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.len)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ u32 │
+    #   # ╞═════╡
+    #   # │ 2   │
+    #   # │ 2   │
+    #   # └─────┘
+    def len
+      Utils.wrap_expr(_rbexpr.arr_len)
+    end
+
+    # Slice every subarray.
+    #
+    # @param offset [Integer]
+    #   Start index. Negative indexing is supported.
+    # @param length [Integer]
+    #   Length of the slice. If set to `None` (default), the slice is taken to the
+    #   end of the list.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `length` and `offset` must be constant values.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.slice(0, 1))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ a         │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [1]       │
+    #   # │ [4]       │
+    #   # └───────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.slice(0, 1, as_array: true))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────────┐
+    #   # │ a             │
+    #   # │ ---           │
+    #   # │ array[i64, 1] │
+    #   # ╞═══════════════╡
+    #   # │ [1]           │
+    #   # │ [4]           │
+    #   # └───────────────┘
+    def slice(
+      offset,
+      length = nil,
+      as_array: false
+    )
+      offset = Utils.parse_into_expression(offset)
+      length = !length.nil? ? Utils.parse_into_expression(length) : nil
+      Utils.wrap_expr(_rbexpr.arr_slice(offset, length, as_array))
+    end
+
+    # Get the first `n` elements of the sub-arrays.
+    #
+    # @param n [Integer]
+    #   Number of values to return for each sublist.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `n` must be a constant value.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.head(1))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ a         │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [1]       │
+    #   # │ [4]       │
+    #   # └───────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.head(1, as_array: true))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────────┐
+    #   # │ a             │
+    #   # │ ---           │
+    #   # │ array[i64, 1] │
+    #   # ╞═══════════════╡
+    #   # │ [1]           │
+    #   # │ [4]           │
+    #   # └───────────────┘
+    def head(n = 5, as_array: false)
+      slice(0, n, as_array: as_array)
+    end
+
+    # Slice the last `n` values of every sublist.
+    #
+    # @param n [Integer]
+    #   Number of values to return for each sublist.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `n` must be a constant value.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.tail(1))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ a         │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [2]       │
+    #   # │ [3]       │
+    #   # └───────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.tail(1, as_array: true))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────────┐
+    #   # │ a             │
+    #   # │ ---           │
+    #   # │ array[i64, 1] │
+    #   # ╞═══════════════╡
+    #   # │ [2]           │
+    #   # │ [3]           │
+    #   # └───────────────┘
+    def tail(n = 5, as_array: false)
+      n = Utils.parse_into_expression(n)
+      Utils.wrap_expr(_rbexpr.arr_tail(n, as_array))
+    end
+
     # Compute the min values of the sub-arrays.
     #
     # @return [Expr]
@@ -30,7 +205,7 @@ module Polars
     #   # │ 3   │
     #   # └─────┘
     def min
-      Utils.wrap_expr(_rbexpr.array_min)
+      Utils.wrap_expr(_rbexpr.arr_min)
     end
 
     # Compute the max values of the sub-arrays.
@@ -54,7 +229,7 @@ module Polars
     #   # │ 4   │
     #   # └─────┘
     def max
-      Utils.wrap_expr(_rbexpr.array_max)
+      Utils.wrap_expr(_rbexpr.arr_max)
     end
 
     # Compute the sum values of the sub-arrays.
@@ -78,7 +253,103 @@ module Polars
     #   # │ 7   │
     #   # └─────┘
     def sum
-      Utils.wrap_expr(_rbexpr.array_sum)
+      Utils.wrap_expr(_rbexpr.arr_sum)
+    end
+
+    # Compute the std of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.std)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌──────────┐
+    #   # │ a        │
+    #   # │ ---      │
+    #   # │ f64      │
+    #   # ╞══════════╡
+    #   # │ 0.707107 │
+    #   # │ 0.707107 │
+    #   # └──────────┘
+    def std(ddof: 1)
+      Utils.wrap_expr(_rbexpr.arr_std(ddof))
+    end
+
+    # Compute the var of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.var)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 0.5 │
+    #   # │ 0.5 │
+    #   # └─────┘
+    def var(ddof: 1)
+      Utils.wrap_expr(_rbexpr.arr_var(ddof))
+    end
+
+    # Compute the mean of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2, 3], [1, 1, 16]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 3)}
+    #   )
+    #   df.select(Polars.col("a").arr.mean)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 2.0 │
+    #   # │ 6.0 │
+    #   # └─────┘
+    def mean
+      Utils.wrap_expr(_rbexpr.arr_mean)
+    end
+
+    # Compute the median of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.median)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 1.5 │
+    #   # │ 3.5 │
+    #   # └─────┘
+    def median
+      Utils.wrap_expr(_rbexpr.arr_median)
     end
 
     # Get the unique/distinct values in the array.
@@ -109,6 +380,32 @@ module Polars
       Utils.wrap_expr(_rbexpr.arr_unique(maintain_order))
     end
 
+    # Count the number of unique values in every sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [[1, 1, 2], [2, 3, 4]],
+    #     },
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 3)}
+    #   )
+    #   df.with_columns(n_unique: Polars.col("a").arr.n_unique)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌───────────────┬──────────┐
+    #   # │ a             ┆ n_unique │
+    #   # │ ---           ┆ ---      │
+    #   # │ array[i64, 3] ┆ u32      │
+    #   # ╞═══════════════╪══════════╡
+    #   # │ [1, 1, 2]     ┆ 2        │
+    #   # │ [2, 3, 4]     ┆ 3        │
+    #   # └───────────────┴──────────┘
+    def n_unique
+      Utils.wrap_expr(_rbexpr.arr_n_unique)
+    end
+
     # Convert an Array column into a List column with the same inner data type.
     #
     # @return [Expr]
@@ -535,5 +832,87 @@ module Polars
       element = Utils.parse_into_expression(element, str_as_lit: true)
       Utils.wrap_expr(_rbexpr.arr_count_matches(element))
     end
+
+    # Convert the Series of type `Array` to a Series of type `Struct`.
+    #
+    # @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 [Expr]
+    #
+    # @example Convert array to struct with default field name assignment:
+    #   df = Polars::DataFrame.new(
+    #     {"n" => [[0, 1, 2], [3, 4, 5]]}, schema: {"n" => Polars::Array.new(Polars::Int8, 3)}
+    #   )
+    #   df.with_columns(struct: Polars.col("n").arr.to_struct)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌──────────────┬───────────┐
+    #   # │ n            ┆ struct    │
+    #   # │ ---          ┆ ---       │
+    #   # │ array[i8, 3] ┆ struct[3] │
+    #   # ╞══════════════╪═══════════╡
+    #   # │ [0, 1, 2]    ┆ {0,1,2}   │
+    #   # │ [3, 4, 5]    ┆ {3,4,5}   │
+    #   # └──────────────┴───────────┘
+    def to_struct(fields: nil)
+      raise Todo if fields
+      if fields.is_a?(Enumerable)
+        field_names = fields.to_a
+        rbexpr = _rbexpr.arr_to_struct(nil)
+        Utils.wrap_expr(rbexpr).struct.rename_fields(field_names)
+      else
+        rbexpr = _rbexpr.arr_to_struct(fields)
+        Utils.wrap_expr(rbexpr)
+      end
+    end
+
+    # Shift array values by the given number of indices.
+    #
+    # @param n [Integer]
+    #   Number of indices to shift forward. If a negative value is passed, values
+    #   are shifted in the opposite direction instead.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   This method is similar to the `LAG` operation in SQL when the value for `n`
+    #   is positive. With a negative value for `n`, it is similar to `LEAD`.
+    #
+    # @example By default, array values are shifted forward by one index.
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2, 3], [4, 5, 6]]}, schema: {"a" => Polars::Array.new(Polars::Int64, 3)}
+    #   )
+    #   df.with_columns(shift: Polars.col("a").arr.shift)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌───────────────┬───────────────┐
+    #   # │ a             ┆ shift         │
+    #   # │ ---           ┆ ---           │
+    #   # │ array[i64, 3] ┆ array[i64, 3] │
+    #   # ╞═══════════════╪═══════════════╡
+    #   # │ [1, 2, 3]     ┆ [null, 1, 2]  │
+    #   # │ [4, 5, 6]     ┆ [null, 4, 5]  │
+    #   # └───────────────┴───────────────┘
+    #
+    # @example Pass a negative value to shift in the opposite direction instead.
+    #   df.with_columns(shift: Polars.col("a").arr.shift(-2))
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌───────────────┬─────────────────┐
+    #   # │ a             ┆ shift           │
+    #   # │ ---           ┆ ---             │
+    #   # │ array[i64, 3] ┆ array[i64, 3]   │
+    #   # ╞═══════════════╪═════════════════╡
+    #   # │ [1, 2, 3]     ┆ [3, null, null] │
+    #   # │ [4, 5, 6]     ┆ [6, null, null] │
+    #   # └───────────────┴─────────────────┘
+    def shift(n = 1)
+      n = Utils.parse_into_expression(n)
+      Utils.wrap_expr(_rbexpr.arr_shift(n))
+    end
   end
 end