Как расширить API Polars для работы как с DataFrame, так и с LazyFrame ⇐ Python

[Anonymous]

Я хотел бы расширить API Polars, как описано в документации, с помощью единого пространства имен, которое должно работать как с DataFrame, так и с LazyFrame.
Представьте, что у меня есть простой DataFrame, который с таким же успехом может быть LazyFrame...

Код: Выделить всё

import polars as pl
a = pl.DataFrame({"col1": ["A", "B", "A", "A", "B"], "col2": [1, 2, 3, 4, 5]})
b = a.lazy()

Я хочу использовать пространство имен, чтобы делать что-то с обоими из них. Для простоты давайте просто создадим простую сумму group_by. Я мог бы написать пространства имен для каждого отдельно:

Код: Выделить всё

@pl.api.register_dataframe_namespace("stuff")
class StuffFrame:
def __init__(self, df: pl.DataFrame):
self._df = df

def do_stuff(self) -> pl.DataFrame:
df = self._df.group_by("col1").agg(pl.col("*").sum())
return df

@pl.api.register_lazyframe_namespace("stuff")
class StuffLazyFrame:
def __init__(self, ldf: pl.LazyFrame):
self._ldf = ldf

def do_stuff(self) -> pl.LazyFrame:
ldf = self._ldf.group_by("col1").agg(pl.col("*").sum())
return ldf

a1 = a.stuff.do_stuff()
b1 = b.stuff.do_stuff()

Но тогда мне придется поддерживать код в двух функциях. Вместо этого я мог бы преобразовать DataFrame в LazyFrame в пространстве имен DataFrame, а затем вызвать пространство имен LazyFrame:

Код: Выделить всё

@pl.api.register_dataframe_namespace("stuff")
class StuffFrame:
def __init__(self, df: pl.DataFrame):
self._df = df

def do_stuff(self) -> pl.DataFrame:
df = self._df.lazy().stuff.do_stuff().collect()  # pl.LazyFrame:
ldf = self._ldf.group_by("col1").agg(pl.col("*").sum()) #

Подробнее здесь: [url]https://stackoverflow.com/questions/78561121/how-to-extend-polars-api-to-work-on-both-dataframe-and-lazyframe[/url]