Расширение API

Добавление новых возможностей

Эти функции позволяют регистрировать пользовательские возможности в отдельном пространстве имён базовых классов Polars без необходимости создавать подклассы или использовать примеси. В качестве целей поддерживаются Expr, DataFrame, LazyFrame и Series.

Эта функция предназначена прежде всего для авторов библиотек, предоставляющих специализированные возможности, которые могут отсутствовать в основной библиотеке (или не относиться к её задачам).

Доступные регистрации

`register_expr_namespace`(name)	Декоратор для регистрации пользовательских возможностей для Expr Polars.
`register_dataframe_namespace`(name)	Декоратор для регистрации пользовательских возможностей для DataFrame Polars.
`register_lazyframe_namespace`(name)	Декоратор для регистрации пользовательских возможностей для LazyFrame Polars.
`register_series_namespace`(name)	Декоратор для регистрации пользовательских возможностей для Series polars.

Примечание

Нельзя переопределять существующие пространства имён Polars (например, .str или .dt); попытка сделать это приведёт к ошибке AttributeError. Однако вы можете переопределять другие пользовательские пространства имён (при этом будет выдано только предупреждение UserWarning).

Примеры

Expr

@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!! │
# └──────────┘    └───────────────┴──────────────────┘

DataFrame

@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 │
#  └─────┘, └─────┘]

LazyFrame

@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 │
# └─────┴─────┴─────┘    └─────┴─────┴─────┘

Series

@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
# ]                    ]                    ]

© 2020 Ritchie Vink
© 2022 Polars contributors
Licensed under the MIT License.
https://docs.pola.rs/api/python/stable/reference/api.html