扩展API#

提供新功能#

这些函数允许您在底层的Polars类中注册自定义功能，而无需子类化或混入。Expr、DataFrame、LazyFrame和Series都是支持的目标。

此功能主要供库作者使用，提供可能不存在（或不属于）核心库的领域特定功能。

可用注册#

register_expr_namespace(name)

用于向Polars Expr注册自定义功能的装饰器。

register_dataframe_namespace(name)

用于向Polars DataFrame注册自定义功能的装饰器。

register_lazyframe_namespace(name)

用于向Polars LazyFrame注册自定义功能的装饰器。

register_series_namespace(name)

用于向polars Series注册自定义功能的装饰器。

注意

你不能覆盖现有的Polars命名空间（如.str或.dt），尝试这样做将会引发一个AttributeError。 然而，你可以覆盖其他自定义命名空间（这只会生成一个UserWarning）。

示例#

@pl.api.register_expr_namespace("greetings")
class Greetings:
    def __init__(self, expr: pl.Expr) -> None:
        self._expr = expr

    def hello(self) -> pl.Expr:
        return (pl.lit("Hello ") + self._expr).alias("hi there")

    def goodbye(self) -> pl.Expr:
        return (pl.lit("Sayōnara ") + self._expr).alias("bye")


pl.DataFrame(data=["world", "world!", "world!!"]).select(
    [
        pl.all().greetings.hello(),
        pl.all().greetings.goodbye(),
    ]
)

# shape: (3, 1)   shape: (3, 2)
# ┌──────────┐    ┌───────────────┬──────────────────┐
# │ column_0 │    │ hi there      ┆ bye              │
# │ ---      │    │ ---           ┆ ---              │
# │ str      │    │ str           ┆ str              │
# ╞══════════╡ >> ╞═══════════════╪══════════════════╡
# │ world    │    │ Hello world   ┆ Sayōnara world   │
# │ world!   │    │ Hello world!  ┆ Sayōnara world!  │
# │ world!!  │    │ Hello world!! ┆ Sayōnara world!! │
# └──────────┘    └───────────────┴──────────────────┘

@pl.api.register_dataframe_namespace("split")
class SplitFrame:
    def __init__(self, df: pl.DataFrame) -> None:
        self._df = df

    def by_alternate_rows(self) -> list[pl.DataFrame]:
        df = self._df.with_row_index(name="n")
        return [
            df.filter((pl.col("n") % 2) == 0).drop("n"),
            df.filter((pl.col("n") % 2) != 0).drop("n"),
        ]


pl.DataFrame(
    data=["aaa", "bbb", "ccc", "ddd", "eee", "fff"],
    schema=[("txt", pl.String)],
).split.by_alternate_rows()

# [┌─────┐  ┌─────┐
#  │ txt │  │ txt │
#  │ --- │  │ --- │
#  │ str │  │ str │
#  ╞═════╡  ╞═════╡
#  │ aaa │  │ bbb │
#  │ ccc │  │ ddd │
#  │ eee │  │ fff │
#  └─────┘, └─────┘]

@pl.api.register_lazyframe_namespace("types")
class DTypeOperations:
    def __init__(self, ldf: pl.LazyFrame) -> None:
        self._ldf = ldf

    def upcast_integer_types(self) -> pl.LazyFrame:
        return self._ldf.with_columns(
            pl.col(tp).cast(pl.Int64)
            for tp in (pl.Int8, pl.Int16, pl.Int32)
        )


ldf = pl.DataFrame(
    data={"a": [1, 2], "b": [3, 4], "c": [5.6, 6.7]},
    schema=[("a", pl.Int16), ("b", pl.Int32), ("c", pl.Float32)],
).lazy()

ldf.types.upcast_integer_types()

# shape: (2, 3)          shape: (2, 3)
# ┌─────┬─────┬─────┐    ┌─────┬─────┬─────┐
# │ a   ┆ b   ┆ c   │    │ a   ┆ b   ┆ c   │
# │ --- ┆ --- ┆ --- │    │ --- ┆ --- ┆ --- │
# │ i16 ┆ i32 ┆ f32 │ >> │ i64 ┆ i64 ┆ f32 │
# ╞═════╪═════╪═════╡    ╞═════╪═════╪═════╡
# │ 1   ┆ 3   ┆ 5.6 │    │ 1   ┆ 3   ┆ 5.6 │
# │ 2   ┆ 4   ┆ 6.7 │    │ 2   ┆ 4   ┆ 6.7 │
# └─────┴─────┴─────┘    └─────┴─────┴─────┘

@pl.api.register_series_namespace("math")
class MathShortcuts:
    def __init__(self, s: pl.Series) -> None:
        self._s = s

    def square(self) -> pl.Series:
        return self._s * self._s

    def cube(self) -> pl.Series:
        return self._s * self._s * self._s


s = pl.Series("n", [1, 2, 3, 4, 5])

s2 = s.math.square().rename("n2")
s3 = s.math.cube().rename("n3")

# shape: (5,)          shape: (5,)           shape: (5,)
# Series: 'n' [i64]    Series: 'n2' [i64]    Series: 'n3' [i64]
# [                    [                     [
#     1                    1                      1
#     2                    4                      8
#     3                    9                      27
#     4                    16                     64
#     5                    25                    125
# ]                    ]                    ]