polars-runtime-compat 1.34.0b2__cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

polars/expr/struct.py

@@ -0,0 +1,357 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+from polars._utils.parse import parse_into_list_of_expressions
+from polars._utils.various import qualified_type_name
+from polars._utils.wrap import wrap_expr
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Sequence
+
+    from polars import Expr
+    from polars._typing import IntoExpr
+
+
+class ExprStructNameSpace:
+    """Namespace for struct related expressions."""
+
+    _accessor = "struct"
+
+    def __init__(self, expr: Expr) -> None:
+        self._pyexpr = expr._pyexpr
+
+    def __getitem__(self, item: str | int) -> Expr:
+        if isinstance(item, str):
+            return self.field(item)
+        elif isinstance(item, int):
+            return wrap_expr(self._pyexpr.struct_field_by_index(item))
+        else:
+            msg = f"expected type 'int | str', got {qualified_type_name(item)!r} ({item!r})"
+            raise TypeError(msg)
+
+    def field(self, name: str | list[str], *more_names: str) -> Expr:
+        """
+        Retrieve one or multiple `Struct` field(s) as a new Series.
+
+        Parameters
+        ----------
+        name
+            Name of the struct field to retrieve.
+        *more_names
+            Additional struct field names.
+
+        Examples
+        --------
+        >>> df = pl.DataFrame(
+        ...     {
+        ...         "aaa": [1, 2],
+        ...         "bbb": ["ab", "cd"],
+        ...         "ccc": [True, None],
+        ...         "ddd": [[1, 2], [3]],
+        ...     }
+        ... ).select(pl.struct("aaa", "bbb", "ccc", "ddd").alias("struct_col"))
+        >>> df
+        shape: (2, 1)
+        ┌──────────────────────┐
+        │ struct_col           │
+        │ ---                  │
+        │ struct[4]            │
+        ╞══════════════════════╡
+        │ {1,"ab",true,[1, 2]} │
+        │ {2,"cd",null,[3]}    │
+        └──────────────────────┘
+
+        Retrieve struct field(s) as Series:
+
+        >>> df.select(pl.col("struct_col").struct.field("bbb"))
+        shape: (2, 1)
+        ┌─────┐
+        │ bbb │
+        │ --- │
+        │ str │
+        ╞═════╡
+        │ ab  │
+        │ cd  │
+        └─────┘
+
+        >>> df.select(
+        ...     pl.col("struct_col").struct.field("bbb"),
+        ...     pl.col("struct_col").struct.field("ddd"),
+        ... )
+        shape: (2, 2)
+        ┌─────┬───────────┐
+        │ bbb ┆ ddd       │
+        │ --- ┆ ---       │
+        │ str ┆ list[i64] │
+        ╞═════╪═══════════╡
+        │ ab  ┆ [1, 2]    │
+        │ cd  ┆ [3]       │
+        └─────┴───────────┘
+
+        Use wildcard expansion:
+
+        >>> df.select(pl.col("struct_col").struct.field("*"))
+        shape: (2, 4)
+        ┌─────┬─────┬──────┬───────────┐
+        │ aaa ┆ bbb ┆ ccc  ┆ ddd       │
+        │ --- ┆ --- ┆ ---  ┆ ---       │
+        │ i64 ┆ str ┆ bool ┆ list[i64] │
+        ╞═════╪═════╪══════╪═══════════╡
+        │ 1   ┆ ab  ┆ true ┆ [1, 2]    │
+        │ 2   ┆ cd  ┆ null ┆ [3]       │
+        └─────┴─────┴──────┴───────────┘
+
+        Retrieve multiple fields by name:
+
+        >>> df.select(pl.col("struct_col").struct.field("aaa", "bbb"))
+        shape: (2, 2)
+        ┌─────┬─────┐
+        │ aaa ┆ bbb │
+        │ --- ┆ --- │
+        │ i64 ┆ str │
+        ╞═════╪═════╡
+        │ 1   ┆ ab  │
+        │ 2   ┆ cd  │
+        └─────┴─────┘
+
+        Retrieve multiple fields by regex expansion:
+
+        >>> df.select(pl.col("struct_col").struct.field("^a.*|b.*$"))
+        shape: (2, 2)
+        ┌─────┬─────┐
+        │ aaa ┆ bbb │
+        │ --- ┆ --- │
+        │ i64 ┆ str │
+        ╞═════╪═════╡
+        │ 1   ┆ ab  │
+        │ 2   ┆ cd  │
+        └─────┴─────┘
+
+        Notes
+        -----
+        The `struct` namespace has implemented `__getitem__`
+        so you can also access fields by index:
+
+        >>> df.select(pl.col("struct_col").struct[1])
+        shape: (2, 1)
+        ┌─────┐
+        │ bbb │
+        │ --- │
+        │ str │
+        ╞═════╡
+        │ ab  │
+        │ cd  │
+        └─────┘
+        """
+        if more_names:
+            name = [*([name] if isinstance(name, str) else name), *more_names]
+        if isinstance(name, list):
+            return wrap_expr(self._pyexpr.struct_multiple_fields(name))
+
+        return wrap_expr(self._pyexpr.struct_field_by_name(name))
+
+    def unnest(self) -> Expr:
+        """
+        Expand the struct into its individual fields.
+
+        Alias for `Expr.struct.field("*")`.
+
+        >>> df = pl.DataFrame(
+        ...     {
+        ...         "aaa": [1, 2],
+        ...         "bbb": ["ab", "cd"],
+        ...         "ccc": [True, None],
+        ...         "ddd": [[1, 2], [3]],
+        ...     }
+        ... ).select(pl.struct("aaa", "bbb", "ccc", "ddd").alias("struct_col"))
+        >>> df
+        shape: (2, 1)
+        ┌──────────────────────┐
+        │ struct_col           │
+        │ ---                  │
+        │ struct[4]            │
+        ╞══════════════════════╡
+        │ {1,"ab",true,[1, 2]} │
+        │ {2,"cd",null,[3]}    │
+        └──────────────────────┘
+        >>> df.select(pl.col("struct_col").struct.unnest())
+        shape: (2, 4)
+        ┌─────┬─────┬──────┬───────────┐
+        │ aaa ┆ bbb ┆ ccc  ┆ ddd       │
+        │ --- ┆ --- ┆ ---  ┆ ---       │
+        │ i64 ┆ str ┆ bool ┆ list[i64] │
+        ╞═════╪═════╪══════╪═══════════╡
+        │ 1   ┆ ab  ┆ true ┆ [1, 2]    │
+        │ 2   ┆ cd  ┆ null ┆ [3]       │
+        └─────┴─────┴──────┴───────────┘
+        """
+        return self.field("*")
+
+    def rename_fields(self, names: Sequence[str]) -> Expr:
+        """
+        Rename the fields of the struct.
+
+        Parameters
+        ----------
+        names
+            New names, given in the same order as the struct's fields.
+
+        Examples
+        --------
+        >>> df = pl.DataFrame(
+        ...     {
+        ...         "aaa": [1, 2],
+        ...         "bbb": ["ab", "cd"],
+        ...         "ccc": [True, None],
+        ...         "ddd": [[1, 2], [3]],
+        ...     }
+        ... ).select(pl.struct("aaa", "bbb", "ccc", "ddd").alias("struct_col"))
+        >>> df
+        shape: (2, 1)
+        ┌──────────────────────┐
+        │ struct_col           │
+        │ ---                  │
+        │ struct[4]            │
+        ╞══════════════════════╡
+        │ {1,"ab",true,[1, 2]} │
+        │ {2,"cd",null,[3]}    │
+        └──────────────────────┘
+
+        >>> df.unnest("struct_col")
+        shape: (2, 4)
+        ┌─────┬─────┬──────┬───────────┐
+        │ aaa ┆ bbb ┆ ccc  ┆ ddd       │
+        │ --- ┆ --- ┆ ---  ┆ ---       │
+        │ i64 ┆ str ┆ bool ┆ list[i64] │
+        ╞═════╪═════╪══════╪═══════════╡
+        │ 1   ┆ ab  ┆ true ┆ [1, 2]    │
+        │ 2   ┆ cd  ┆ null ┆ [3]       │
+        └─────┴─────┴──────┴───────────┘
+
+        Rename fields:
+
+        >>> df = df.select(
+        ...     pl.col("struct_col").struct.rename_fields(["www", "xxx", "yyy", "zzz"])
+        ... )
+        >>> df.unnest("struct_col")
+        shape: (2, 4)
+        ┌─────┬─────┬──────┬───────────┐
+        │ www ┆ xxx ┆ yyy  ┆ zzz       │
+        │ --- ┆ --- ┆ ---  ┆ ---       │
+        │ i64 ┆ str ┆ bool ┆ list[i64] │
+        ╞═════╪═════╪══════╪═══════════╡
+        │ 1   ┆ ab  ┆ true ┆ [1, 2]    │
+        │ 2   ┆ cd  ┆ null ┆ [3]       │
+        └─────┴─────┴──────┴───────────┘
+
+        Following a rename, the previous field names (obviously) cannot be referenced:
+
+        >>> df.select(pl.col("struct_col").struct.field("aaa"))  # doctest: +SKIP
+        StructFieldNotFoundError: aaa
+        """
+        return wrap_expr(self._pyexpr.struct_rename_fields(names))
+
+    def json_encode(self) -> Expr:
+        """
+        Convert this struct to a string column with json values.
+
+        Examples
+        --------
+        >>> pl.DataFrame(
+        ...     {"a": [{"a": [1, 2], "b": [45]}, {"a": [9, 1, 3], "b": None}]}
+        ... ).with_columns(pl.col("a").struct.json_encode().alias("encoded"))
+        shape: (2, 2)
+        ┌──────────────────┬────────────────────────┐
+        │ a                ┆ encoded                │
+        │ ---              ┆ ---                    │
+        │ struct[2]        ┆ str                    │
+        ╞══════════════════╪════════════════════════╡
+        │ {[1, 2],[45]}    ┆ {"a":[1,2],"b":[45]}   │
+        │ {[9, 1, 3],null} ┆ {"a":[9,1,3],"b":null} │
+        └──────────────────┴────────────────────────┘
+        """
+        return wrap_expr(self._pyexpr.struct_json_encode())
+
+    def with_fields(
+        self,
+        *exprs: IntoExpr | Iterable[IntoExpr],
+        **named_exprs: IntoExpr,
+    ) -> Expr:
+        """
+        Add or overwrite fields of this struct.
+
+        This is similar to `with_columns` on `DataFrame`.
+
+        .. versionadded:: 0.20.27
+
+        Examples
+        --------
+        >>> df = pl.DataFrame(
+        ...     {
+        ...         "coords": [{"x": 1, "y": 4}, {"x": 4, "y": 9}, {"x": 9, "y": 16}],
+        ...         "multiply": [10, 2, 3],
+        ...     }
+        ... )
+        >>> df
+        shape: (3, 2)
+        ┌───────────┬──────────┐
+        │ coords    ┆ multiply │
+        │ ---       ┆ ---      │
+        │ struct[2] ┆ i64      │
+        ╞═══════════╪══════════╡
+        │ {1,4}     ┆ 10       │
+        │ {4,9}     ┆ 2        │
+        │ {9,16}    ┆ 3        │
+        └───────────┴──────────┘
+        >>> df = df.with_columns(
+        ...     pl.col("coords").struct.with_fields(
+        ...         pl.field("x").sqrt(),
+        ...         y_mul=pl.field("y") * pl.col("multiply"),
+        ...     )
+        ... )
+        >>> df
+        shape: (3, 2)
+        ┌─────────────┬──────────┐
+        │ coords      ┆ multiply │
+        │ ---         ┆ ---      │
+        │ struct[3]   ┆ i64      │
+        ╞═════════════╪══════════╡
+        │ {1.0,4,40}  ┆ 10       │
+        │ {2.0,9,18}  ┆ 2        │
+        │ {3.0,16,48} ┆ 3        │
+        └─────────────┴──────────┘
+        >>> df.unnest("coords")
+        shape: (3, 4)
+        ┌─────┬─────┬───────┬──────────┐
+        │ x   ┆ y   ┆ y_mul ┆ multiply │
+        │ --- ┆ --- ┆ ---   ┆ ---      │
+        │ f64 ┆ i64 ┆ i64   ┆ i64      │
+        ╞═════╪═════╪═══════╪══════════╡
+        │ 1.0 ┆ 4   ┆ 40    ┆ 10       │
+        │ 2.0 ┆ 9   ┆ 18    ┆ 2        │
+        │ 3.0 ┆ 16  ┆ 48    ┆ 3        │
+        └─────┴─────┴───────┴──────────┘
+
+        Parameters
+        ----------
+        *exprs
+            Field(s) to add, specified as positional arguments.
+            Accepts expression input. Strings are parsed as column names, other
+            non-expression inputs are parsed as literals.
+        **named_exprs
+            Additional fields to add, specified as keyword arguments.
+            The columns will be renamed to the keyword used.
+
+        See Also
+        --------
+        field
+        """
+        structify = bool(int(os.environ.get("POLARS_AUTO_STRUCTIFY", 0)))
+
+        pyexprs = parse_into_list_of_expressions(
+            *exprs, **named_exprs, __structify=structify
+        )
+
+        return wrap_expr(self._pyexpr.struct_with_fields(pyexprs))

polars/expr/whenthen.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import polars.functions as F
+from polars._utils.parse import (
+    parse_into_expression,
+    parse_predicates_constraints_into_expression,
+)
+from polars._utils.wrap import wrap_expr
+from polars.expr.expr import Expr
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from polars._plr import PyExpr
+    from polars._typing import IntoExpr
+
+
+class When:
+    """
+    Utility class for the `when-then-otherwise` expression.
+
+    Represents the initial state of the expression after `pl.when(...)` is called.
+
+    In this state, `then` must be called to continue to finish the expression.
+    """
+
+    def __init__(self, when: Any) -> None:
+        self._when = when
+
+    def then(self, statement: IntoExpr) -> Then:
+        """
+        Attach a statement to the corresponding condition.
+
+        Parameters
+        ----------
+        statement
+            The statement to apply if the corresponding condition is true.
+            Accepts expression input. Strings are parsed as column names, other
+            non-expression inputs are parsed as literals.
+        """
+        statement_pyexpr = parse_into_expression(statement)
+        return Then(self._when.then(statement_pyexpr))
+
+
+class Then(Expr):
+    """
+    Utility class for the `when-then-otherwise` expression.
+
+    Represents the state of the expression after `pl.when(...).then(...)` is called.
+    """
+
+    def __init__(self, then: Any) -> None:
+        self._then = then
+
+    @classmethod
+    def _from_pyexpr(cls, pyexpr: PyExpr) -> Expr:
+        return wrap_expr(pyexpr)
+
+    @property
+    def _pyexpr(self) -> PyExpr:  # type: ignore[override]
+        return self._then.otherwise(F.lit(None)._pyexpr)
+
+    def when(
+        self,
+        *predicates: IntoExpr | Iterable[IntoExpr],
+        **constraints: Any,
+    ) -> ChainedWhen:
+        """
+        Add a condition to the `when-then-otherwise` expression.
+
+        Parameters
+        ----------
+        predicates
+            Condition(s) that must be met in order to apply the subsequent statement.
+            Accepts one or more boolean expressions, which are implicitly combined with
+            `&`. String input is parsed as a column name.
+        constraints
+            Apply conditions as `col_name = value` keyword arguments that are treated as
+            equality matches, such as `x = 123`. As with the predicates parameter,
+            multiple conditions are implicitly combined using `&`.
+        """
+        condition_pyexpr = parse_predicates_constraints_into_expression(
+            *predicates, **constraints
+        )
+        return ChainedWhen(self._then.when(condition_pyexpr))
+
+    def otherwise(self, statement: IntoExpr) -> Expr:
+        """
+        Define a default for the `when-then-otherwise` expression.
+
+        Parameters
+        ----------
+        statement
+            The statement to apply if all conditions are false.
+            Accepts expression input. Strings are parsed as column names, other
+            non-expression inputs are parsed as literals.
+        """
+        statement_pyexpr = parse_into_expression(statement)
+        return wrap_expr(self._then.otherwise(statement_pyexpr))
+
+
+class ChainedWhen:
+    """
+    Utility class for the `when-then-otherwise` expression.
+
+    Represents the state of the expression after an additional `when` is called.
+
+    In this state, `then` must be called to continue to finish the expression.
+    """
+
+    def __init__(self, chained_when: Any) -> None:
+        self._chained_when = chained_when
+
+    def then(self, statement: IntoExpr) -> ChainedThen:
+        """
+        Attach a statement to the corresponding condition.
+
+        Parameters
+        ----------
+        statement
+            The statement to apply if the corresponding condition is true.
+            Accepts expression input. Strings are parsed as column names, other
+            non-expression inputs are parsed as literals.
+        """
+        statement_pyexpr = parse_into_expression(statement)
+        return ChainedThen(self._chained_when.then(statement_pyexpr))
+
+
+class ChainedThen(Expr):
+    """
+    Utility class for the `when-then-otherwise` expression.
+
+    Represents the state of the expression after an additional `then` is called.
+    """
+
+    def __init__(self, chained_then: Any) -> None:
+        self._chained_then = chained_then
+
+    @classmethod
+    def _from_pyexpr(cls, pyexpr: PyExpr) -> Expr:
+        return wrap_expr(pyexpr)
+
+    @property
+    def _pyexpr(self) -> PyExpr:  # type: ignore[override]
+        return self._chained_then.otherwise(F.lit(None)._pyexpr)
+
+    def when(
+        self,
+        *predicates: IntoExpr | Iterable[IntoExpr],
+        **constraints: Any,
+    ) -> ChainedWhen:
+        """
+        Add another condition to the `when-then-otherwise` expression.
+
+        Parameters
+        ----------
+        predicates
+            Condition(s) that must be met in order to apply the subsequent statement.
+            Accepts one or more boolean expressions, which are implicitly combined with
+            `&`. String input is parsed as a column name.
+        constraints
+            Apply conditions as `col_name = value` keyword arguments that are treated as
+            equality matches, such as `x = 123`. As with the predicates parameter,
+            multiple conditions are implicitly combined using `&`.
+        """
+        condition_pyexpr = parse_predicates_constraints_into_expression(
+            *predicates, **constraints
+        )
+        return ChainedWhen(self._chained_then.when(condition_pyexpr))
+
+    def otherwise(self, statement: IntoExpr) -> Expr:
+        """
+        Define a default for the `when-then-otherwise` expression.
+
+        Parameters
+        ----------
+        statement
+            The statement to apply if all conditions are false.
+            Accepts expression input. Strings are parsed as column names, other
+            non-expression inputs are parsed as literals.
+        """
+        statement_pyexpr = parse_into_expression(statement)
+        return wrap_expr(self._chained_then.otherwise(statement_pyexpr))