Ejemplo n.º 1
0
def fold(
    acc: "pl.Expr",
    f: Callable[["pl.Series", "pl.Series"], "pl.Series"],
    exprs: Union[tp.List["pl.Expr"], "pl.Expr"],
) -> "pl.Expr":
    """
    Accumulate over multiple columns horizontally/ row wise with a left fold.

    Parameters
    ----------
    acc
     Accumulator Expression. This is the value that will be initialized when the fold starts.
     For a sum this could for instance be lit(0).

    f
        Function to apply over the accumulator and the value.
        Fn(acc, value) -> new_value
    exprs
        Expressions to aggregate over. May also be a wildcard expression.
    """
    # in case of pl.col("*")
    acc = pl.lazy.expr.expr_to_lit_or_expr(acc, str_to_lit=True)
    if isinstance(exprs, pl.Expr):
        exprs = [exprs]

    exprs = pl.lazy.expr._selection_to_pyexpr_list(exprs)
    return pl.wrap_expr(pyfold(acc._pyexpr, f, exprs))
Ejemplo n.º 2
0
def arange(
    low: Union[int, "pl.Expr", "pl.Series"],
    high: Union[int, "pl.Expr", "pl.Series"],
    step: int = 1,
    dtype: Optional[Type[DataType]] = None,
    eager: bool = False,
) -> Union["pl.Expr", "pl.Series"]:
    """
    Create a range expression. This can be used in a `select`, `with_column` etc.
    Be sure that the range size is equal to the DataFrame you are collecting.

     Examples
     --------

    >>> (df.lazy()
        .filter(pl.col("foo") < pl.arange(0, 100))
        .collect())

    Parameters
    ----------
    low
        Lower bound of range.
    high
        Upper bound of range.
    step
        Step size of the range
    dtype
        deprecated, cast later
    eager
        If eager evaluation is `True`, a Series is returned instead of an Expr
    """
    low = pl.lazy.expr_to_lit_or_expr(low, str_to_lit=False)
    high = pl.lazy.expr_to_lit_or_expr(high, str_to_lit=False)

    if eager:
        df = pl.DataFrame({"a": [1]})
        return df.select(pl.arange(
            low, high, step).alias("arange"))["arange"]  # type: ignore

    return pl.wrap_expr(pyarange(low._pyexpr, high._pyexpr, step))