polars-runtime-compat 1.34.0b2__cp39-abi3-manylinux_2_24_aarch64.whl

polars/functions/repeat.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+import contextlib
+from decimal import Decimal as D
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, overload
+
+from polars import functions as F
+from polars._utils.parse import parse_into_expression
+from polars._utils.various import qualified_type_name
+from polars._utils.wrap import wrap_expr
+from polars.datatypes import (
+    Array,
+    Boolean,
+    Decimal,
+    Float64,
+    List,
+    Utf8,
+)
+from polars.datatypes.group import FLOAT_DTYPES, INTEGER_DTYPES
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    import polars._plr as plr
+
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+    from polars import Expr, Series
+    from polars._typing import IntoExpr, PolarsDataType
+
+
+# create a lookup of dtypes that have a reasonable one/zero mapping; for
+# anything more elaborate should use `repeat`
+@lru_cache(16)
+def _one_or_zero_by_dtype(value: int, dtype: PolarsDataType) -> Any:
+    if dtype in INTEGER_DTYPES:
+        return value
+    elif dtype in FLOAT_DTYPES:
+        return float(value)
+    elif dtype == Boolean:
+        return bool(value)
+    elif dtype == Utf8:
+        return str(value)
+    elif isinstance(dtype, Decimal):
+        return D(value)
+    elif isinstance(dtype, (List, Array)):
+        arr_width = getattr(dtype, "size", 1)
+        return [_one_or_zero_by_dtype(value, dtype.inner)] * arr_width
+    return None
+
+
+@overload
+def repeat(
+    value: IntoExpr | None,
+    n: int | Expr,
+    *,
+    dtype: PolarsDataType | None = ...,
+    eager: Literal[False] = ...,
+) -> Expr: ...
+
+
+@overload
+def repeat(
+    value: IntoExpr | None,
+    n: int | Expr,
+    *,
+    dtype: PolarsDataType | None = ...,
+    eager: Literal[True],
+) -> Series: ...
+
+
+@overload
+def repeat(
+    value: IntoExpr | None,
+    n: int | Expr,
+    *,
+    dtype: PolarsDataType | None = ...,
+    eager: bool,
+) -> Expr | Series: ...
+
+
+def repeat(
+    value: IntoExpr | None,
+    n: int | Expr,
+    *,
+    dtype: PolarsDataType | None = None,
+    eager: bool = False,
+) -> Expr | Series:
+    """
+    Construct a column of length `n` filled with the given value.
+
+    Parameters
+    ----------
+    value
+        Value to repeat.
+    n
+        Length of the resulting column.
+    dtype
+        Data type of the resulting column. If set to `None` (default), data type is
+        inferred from the given value. Defaults to Int32 for integer values, unless
+        Int64 is required to fit the given value. Defaults to Float64 for float values.
+    eager
+        Evaluate immediately and return a `Series`. If set to `False` (default),
+        return an expression instead.
+
+    Notes
+    -----
+    If you want to construct a column in lazy mode and do not need a pre-determined
+    length, use :func:`lit` instead.
+
+    See Also
+    --------
+    lit
+
+    Examples
+    --------
+    Construct a column with a repeated value in a lazy context.
+
+    >>> pl.select(pl.repeat("z", n=3)).to_series()
+    shape: (3,)
+    Series: 'repeat' [str]
+    [
+            "z"
+            "z"
+            "z"
+    ]
+
+    Generate a Series directly by setting `eager=True`.
+
+    >>> pl.repeat(3, n=3, dtype=pl.Int8, eager=True)
+    shape: (3,)
+    Series: 'repeat' [i8]
+    [
+            3
+            3
+            3
+    ]
+    """
+    if isinstance(n, int):
+        n = F.lit(n)
+    if not hasattr(n, "_pyexpr"):
+        msg = f"`n` parameter of `repeat expected a `int` or `Expr` got a `{qualified_type_name(n)}`"
+        raise TypeError(msg)
+    value_pyexpr = parse_into_expression(value, str_as_lit=True, dtype=dtype)
+    expr = wrap_expr(plr.repeat(value_pyexpr, n._pyexpr, dtype))
+    if eager:
+        return F.select(expr).to_series()
+    return expr
+
+
+@overload
+def ones(
+    n: int | Expr,
+    dtype: PolarsDataType = ...,
+    *,
+    eager: Literal[False] = ...,
+) -> Expr: ...
+
+
+@overload
+def ones(
+    n: int | Expr,
+    dtype: PolarsDataType = ...,
+    *,
+    eager: Literal[True],
+) -> Series: ...
+
+
+@overload
+def ones(
+    n: int | Expr,
+    dtype: PolarsDataType = ...,
+    *,
+    eager: bool,
+) -> Expr | Series: ...
+
+
+def ones(
+    n: int | Expr,
+    dtype: PolarsDataType = Float64,
+    *,
+    eager: bool = False,
+) -> Expr | Series:
+    """
+    Construct a column of length `n` filled with ones.
+
+    This is syntactic sugar for the `repeat` function.
+
+    Parameters
+    ----------
+    n
+        Length of the resulting column.
+    dtype
+        Data type of the resulting column. Defaults to Float64.
+    eager
+        Evaluate immediately and return a `Series`. If set to `False`,
+        return an expression instead.
+
+    Notes
+    -----
+    If you want to construct a column in lazy mode and do not need a pre-determined
+    length, use :func:`lit` instead.
+
+    See Also
+    --------
+    repeat
+    lit
+
+    Examples
+    --------
+    >>> pl.ones(3, pl.Int8, eager=True)
+    shape: (3,)
+    Series: 'ones' [i8]
+    [
+        1
+        1
+        1
+    ]
+    """
+    if (one := _one_or_zero_by_dtype(1, dtype)) is None:
+        msg = f"invalid dtype for `ones`; found {dtype}"
+        raise TypeError(msg)
+
+    return repeat(one, n=n, dtype=dtype, eager=eager).alias("ones")
+
+
+@overload
+def zeros(
+    n: int | Expr,
+    dtype: PolarsDataType = ...,
+    *,
+    eager: Literal[False] = ...,
+) -> Expr: ...
+
+
+@overload
+def zeros(
+    n: int | Expr,
+    dtype: PolarsDataType = ...,
+    *,
+    eager: Literal[True],
+) -> Series: ...
+
+
+@overload
+def zeros(
+    n: int | Expr,
+    dtype: PolarsDataType = ...,
+    *,
+    eager: bool,
+) -> Expr | Series: ...
+
+
+def zeros(
+    n: int | Expr,
+    dtype: PolarsDataType = Float64,
+    *,
+    eager: bool = False,
+) -> Expr | Series:
+    """
+    Construct a column of length `n` filled with zeros.
+
+    This is syntactic sugar for the `repeat` function.
+
+    Parameters
+    ----------
+    n
+        Length of the resulting column.
+    dtype
+        Data type of the resulting column. Defaults to Float64.
+    eager
+        Evaluate immediately and return a `Series`. If set to `False`,
+        return an expression instead.
+
+    Notes
+    -----
+    If you want to construct a column in lazy mode and do not need a pre-determined
+    length, use :func:`lit` instead.
+
+    See Also
+    --------
+    repeat
+    lit
+
+    Examples
+    --------
+    >>> pl.zeros(3, pl.Int8, eager=True)
+    shape: (3,)
+    Series: 'zeros' [i8]
+    [
+        0
+        0
+        0
+    ]
+    """
+    if (zero := _one_or_zero_by_dtype(0, dtype)) is None:
+        msg = f"invalid dtype for `zeros`; found {dtype}"
+        raise TypeError(msg)
+
+    return repeat(zero, n=n, dtype=dtype, eager=eager).alias("zeros")

polars/functions/whenthen.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import contextlib
+from typing import TYPE_CHECKING, Any
+
+import polars._reexport as pl
+from polars._utils.parse import parse_predicates_constraints_into_expression
+
+with contextlib.suppress(ImportError):  # Module not available when building docs
+    import polars._plr as plr
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+    from polars._typing import IntoExprColumn
+
+
+def when(
+    *predicates: IntoExprColumn | Iterable[IntoExprColumn] | bool,
+    **constraints: Any,
+) -> pl.When:
+    """
+    Start a `when-then-otherwise` expression.
+
+    Always initiated by a `pl.when().then()`., and optionally followed by chaining one
+    or more `.when().then()` statements.
+
+    An optional `.otherwise()` can be appended at the end. If not declared, a default
+    of `.otherwise(None)` is used.
+
+    Similar to :func:`coalesce`, the value from the first condition that
+    evaluates to True will be picked.
+
+    If all conditions are False, the `otherwise` value is picked.
+
+    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
+        `&`.
+    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 `&`.
+
+    Warnings
+    --------
+    Polars computes all expressions passed to `when-then-otherwise` in parallel and
+    filters afterwards. This means each expression must be valid on its own, regardless
+    of the conditions in the `when-then-otherwise` chain.
+
+    Notes
+    -----
+    String inputs e.g. `when("string")`, `then("string")` or `otherwise("string")`
+    are parsed as column names. :func:`lit` can be used to create string values.
+
+    Examples
+    --------
+    Below we add a column with the value 1, where column "foo" > 2 and the value
+    1 + column "bar" where it isn't.
+
+    >>> df = pl.DataFrame({"foo": [1, 3, 4], "bar": [3, 4, 0]})
+    >>> df.with_columns(
+    ...     pl.when(pl.col.foo > 2).then(1).otherwise(1 + pl.col.bar).alias("val")
+    ... )
+    shape: (3, 3)
+    ┌─────┬─────┬─────┐
+    │ foo ┆ bar ┆ val │
+    │ --- ┆ --- ┆ --- │
+    │ i64 ┆ i64 ┆ i64 │
+    ╞═════╪═════╪═════╡
+    │ 1   ┆ 3   ┆ 4   │
+    │ 3   ┆ 4   ┆ 1   │
+    │ 4   ┆ 0   ┆ 1   │
+    └─────┴─────┴─────┘
+
+    Note that `when-then` always executes all expressions.
+
+    The results are folded left to right, picking the `then` value from the first `when`
+    condition that is True.
+
+    If no `when` condition is True the `otherwise` value is picked.
+
+    >>> df.with_columns(
+    ...     when = pl.col.foo > 2,
+    ...     then = 1,
+    ...     otherwise = 1 + pl.col.bar
+    ... ).with_columns(
+    ...     pl.when("when").then("then").otherwise("otherwise").alias("val")
+    ... )
+    shape: (3, 6)
+    ┌─────┬─────┬───────┬──────┬───────────┬─────┐
+    │ foo ┆ bar ┆ when  ┆ then ┆ otherwise ┆ val │
+    │ --- ┆ --- ┆ ---   ┆ ---  ┆ ---       ┆ --- │
+    │ i64 ┆ i64 ┆ bool  ┆ i32  ┆ i64       ┆ i64 │
+    ╞═════╪═════╪═══════╪══════╪═══════════╪═════╡
+    │ 1   ┆ 3   ┆ false ┆ 1    ┆ 4         ┆ 4   │
+    │ 3   ┆ 4   ┆ true  ┆ 1    ┆ 5         ┆ 1   │
+    │ 4   ┆ 0   ┆ true  ┆ 1    ┆ 1         ┆ 1   │
+    └─────┴─────┴───────┴──────┴───────────┴─────┘
+
+    Note that in regular Polars usage, a single string is parsed as a column name.
+
+    >>> df.with_columns(
+    ...     when = pl.col.foo > 2,
+    ...     then = "foo",
+    ...     otherwise = "bar"
+    ... )
+    shape: (3, 5)
+    ┌─────┬─────┬───────┬──────┬───────────┐
+    │ foo ┆ bar ┆ when  ┆ then ┆ otherwise │
+    │ --- ┆ --- ┆ ---   ┆ ---  ┆ ---       │
+    │ i64 ┆ i64 ┆ bool  ┆ i64  ┆ i64       │
+    ╞═════╪═════╪═══════╪══════╪═══════════╡
+    │ 1   ┆ 3   ┆ false ┆ 1    ┆ 3         │
+    │ 3   ┆ 4   ┆ true  ┆ 3    ┆ 4         │
+    │ 4   ┆ 0   ┆ true  ┆ 4    ┆ 0         │
+    └─────┴─────┴───────┴──────┴───────────┘
+
+    For consistency, `when-then` behaves in the same way.
+
+    >>> df.with_columns(
+    ...     pl.when(pl.col.foo > 2).then("foo").otherwise("bar").alias("val")
+    ... )
+    shape: (3, 3)
+    ┌─────┬─────┬─────┐
+    │ foo ┆ bar ┆ val │
+    │ --- ┆ --- ┆ --- │
+    │ i64 ┆ i64 ┆ i64 │
+    ╞═════╪═════╪═════╡
+    │ 1   ┆ 3   ┆ 3   │
+    │ 3   ┆ 4   ┆ 3   │
+    │ 4   ┆ 0   ┆ 4   │
+    └─────┴─────┴─────┘
+
+    :func:`lit` can be used to create string values.
+
+    >>> df.with_columns(
+    ...     pl.when(pl.col.foo > 2)
+    ...     .then(pl.lit("foo"))
+    ...     .otherwise(pl.lit("bar"))
+    ...     .alias("val")
+    ... )
+    shape: (3, 3)
+    ┌─────┬─────┬─────┐
+    │ foo ┆ bar ┆ val │
+    │ --- ┆ --- ┆ --- │
+    │ i64 ┆ i64 ┆ str │
+    ╞═════╪═════╪═════╡
+    │ 1   ┆ 3   ┆ bar │
+    │ 3   ┆ 4   ┆ foo │
+    │ 4   ┆ 0   ┆ foo │
+    └─────┴─────┴─────┘
+
+    Multiple `when-then` statements can be chained.
+
+    >>> df.with_columns(
+    ...     pl.when(pl.col.foo > 2)
+    ...     .then(1)
+    ...     .when(pl.col.bar > 2)
+    ...     .then(4)
+    ...     .otherwise(-1)
+    ...     .alias("val")
+    ... )
+    shape: (3, 3)
+    ┌─────┬─────┬─────┐
+    │ foo ┆ bar ┆ val │
+    │ --- ┆ --- ┆ --- │
+    │ i64 ┆ i64 ┆ i32 │
+    ╞═════╪═════╪═════╡
+    │ 1   ┆ 3   ┆ 4   │
+    │ 3   ┆ 4   ┆ 1   │
+    │ 4   ┆ 0   ┆ 1   │
+    └─────┴─────┴─────┘
+
+    In the case of `foo=3` and `bar=4`, both conditions are True but the first value
+    (i.e. 1) is picked.
+
+    >>> df.with_columns(
+    ...     when1 = pl.col.foo > 2,
+    ...     then1 = 1,
+    ...     when2 = pl.col.bar > 2,
+    ...     then2 = 4,
+    ...     otherwise = -1
+    ... )
+    shape: (3, 7)
+    ┌─────┬─────┬───────┬───────┬───────┬───────┬───────────┐
+    │ foo ┆ bar ┆ when1 ┆ then1 ┆ when2 ┆ then2 ┆ otherwise │
+    │ --- ┆ --- ┆ ---   ┆ ---   ┆ ---   ┆ ---   ┆ ---       │
+    │ i64 ┆ i64 ┆ bool  ┆ i32   ┆ bool  ┆ i32   ┆ i32       │
+    ╞═════╪═════╪═══════╪═══════╪═══════╪═══════╪═══════════╡
+    │ 1   ┆ 3   ┆ false ┆ 1     ┆ true  ┆ 4     ┆ -1        │
+    │ 3   ┆ 4   ┆ true  ┆ 1     ┆ true  ┆ 4     ┆ -1        │
+    │ 4   ┆ 0   ┆ true  ┆ 1     ┆ false ┆ 4     ┆ -1        │
+    └─────┴─────┴───────┴───────┴───────┴───────┴───────────┘
+
+    The `otherwise` statement is optional and defaults to `.otherwise(None)`
+    if not given.
+
+    This idiom is commonly used to null out values.
+
+    >>> df.with_columns(pl.when(pl.col.foo == 3).then("bar"))
+    shape: (3, 2)
+    ┌─────┬──────┐
+    │ foo ┆ bar  │
+    │ --- ┆ ---  │
+    │ i64 ┆ i64  │
+    ╞═════╪══════╡
+    │ 1   ┆ null │
+    │ 3   ┆ 4    │
+    │ 4   ┆ null │
+    └─────┴──────┘
+
+    `when` accepts keyword arguments as shorthand for equality conditions.
+
+    >>> df.with_columns(pl.when(foo=3).then("bar"))
+    shape: (3, 2)
+    ┌─────┬──────┐
+    │ foo ┆ bar  │
+    │ --- ┆ ---  │
+    │ i64 ┆ i64  │
+    ╞═════╪══════╡
+    │ 1   ┆ null │
+    │ 3   ┆ 4    │
+    │ 4   ┆ null │
+    └─────┴──────┘
+
+    Multiple predicates passed to `when` are combined with `&`
+
+    >>> df.with_columns(
+    ...     pl.when(pl.col.foo > 2, pl.col.bar < 3) # when((pred1) & (pred2))
+    ...     .then(pl.lit("Yes"))
+    ...     .otherwise(pl.lit("No"))
+    ...     .alias("val")
+    ... )
+    shape: (3, 3)
+    ┌─────┬─────┬─────┐
+    │ foo ┆ bar ┆ val │
+    │ --- ┆ --- ┆ --- │
+    │ i64 ┆ i64 ┆ str │
+    ╞═════╪═════╪═════╡
+    │ 1   ┆ 3   ┆ No  │
+    │ 3   ┆ 4   ┆ No  │
+    │ 4   ┆ 0   ┆ Yes │
+    └─────┴─────┴─────┘
+
+    It could also be thought of as an implicit :func:`all_horizontal` being present.
+
+    >>> df.with_columns(
+    ...     when = pl.all_horizontal(pl.col.foo > 2, pl.col.bar < 3)
+    ... )
+    shape: (3, 3)
+    ┌─────┬─────┬───────┐
+    │ foo ┆ bar ┆ when  │
+    │ --- ┆ --- ┆ ---   │
+    │ i64 ┆ i64 ┆ bool  │
+    ╞═════╪═════╪═══════╡
+    │ 1   ┆ 3   ┆ false │
+    │ 3   ┆ 4   ┆ false │
+    │ 4   ┆ 0   ┆ true  │
+    └─────┴─────┴───────┘
+
+    Structs can be used as a way to return multiple values.
+
+    Here we swap the "foo" and "bar" values when "foo" is greater than 2.
+
+    >>> df.with_columns(
+    ...     pl.when(pl.col.foo > 2)
+    ...     .then(pl.struct(foo="bar", bar="foo"))
+    ...     .otherwise(pl.struct("foo", "bar"))
+    ...     .struct.unnest()
+    ... )
+    shape: (3, 2)
+    ┌─────┬─────┐
+    │ foo ┆ bar │
+    │ --- ┆ --- │
+    │ i64 ┆ i64 │
+    ╞═════╪═════╡
+    │ 1   ┆ 3   │
+    │ 4   ┆ 3   │
+    │ 0   ┆ 4   │
+    └─────┴─────┘
+
+    The struct fields are given the same name as the target columns, which are then
+    unnested.
+
+    >>> df.with_columns(
+    ...     when = pl.col.foo > 2,
+    ...     then = pl.struct(foo="bar", bar="foo"),
+    ...     otherwise = pl.struct("foo", "bar")
+    ... )
+    shape: (3, 5)
+    ┌─────┬─────┬───────┬───────────┬───────────┐
+    │ foo ┆ bar ┆ when  ┆ then      ┆ otherwise │
+    │ --- ┆ --- ┆ ---   ┆ ---       ┆ ---       │
+    │ i64 ┆ i64 ┆ bool  ┆ struct[2] ┆ struct[2] │
+    ╞═════╪═════╪═══════╪═══════════╪═══════════╡
+    │ 1   ┆ 3   ┆ false ┆ {3,1}     ┆ {1,3}     │
+    │ 3   ┆ 4   ┆ true  ┆ {4,3}     ┆ {3,4}     │
+    │ 4   ┆ 0   ┆ true  ┆ {0,4}     ┆ {4,0}     │
+    └─────┴─────┴───────┴───────────┴───────────┘
+
+    The output name of a `when-then` expression comes from the first `then` branch.
+
+    Here we try to set all columns to 0 if any column contains a value less than 2.
+
+    >>> df.with_columns( # doctest: +SKIP
+    ...    pl.when(pl.any_horizontal(pl.all() < 2))
+    ...    .then(0)
+    ...    .otherwise(pl.all())
+    ... )
+    # ComputeError: the name 'literal' passed to `LazyFrame.with_columns` is duplicate
+
+    :meth:`.name.keep` could be used to give preference to the column expression.
+
+    >>> df.with_columns(
+    ...    pl.when(pl.any_horizontal(pl.all() < 2))
+    ...    .then(0)
+    ...    .otherwise(pl.all())
+    ...    .name.keep()
+    ... )
+    shape: (3, 2)
+    ┌─────┬─────┐
+    │ foo ┆ bar │
+    │ --- ┆ --- │
+    │ i64 ┆ i64 │
+    ╞═════╪═════╡
+    │ 0   ┆ 0   │
+    │ 3   ┆ 4   │
+    │ 0   ┆ 0   │
+    └─────┴─────┘
+
+    The logic could also be changed to move the column expression inside `then`.
+
+    >>> df.with_columns(
+    ...     pl.when(pl.any_horizontal(pl.all() < 2).not_())
+    ...     .then(pl.all())
+    ...     .otherwise(0)
+    ... )
+    shape: (3, 2)
+    ┌─────┬─────┐
+    │ foo ┆ bar │
+    │ --- ┆ --- │
+    │ i64 ┆ i64 │
+    ╞═════╪═════╡
+    │ 0   ┆ 0   │
+    │ 3   ┆ 4   │
+    │ 0   ┆ 0   │
+    └─────┴─────┘
+    """  # fmt: skip
+    condition = parse_predicates_constraints_into_expression(*predicates, **constraints)
+    return pl.When(plr.when(condition))

polars/interchange/__init__.py

@@ -0,0 +1,10 @@
+"""
+Module containing the implementation of the Python dataframe interchange protocol.
+
+Details on the protocol:
+https://data-apis.org/dataframe-protocol/latest/index.html
+"""
+
+from polars.interchange.protocol import CompatLevel
+
+__all__ = ["CompatLevel"]