[Python] How to extend Polars API to work on both DataFrame and LazyFrame

I would like to extend the Polars API as described in the docs, but I would like the same namespace to work on both DataFrames and LazyFrames.

Pretend I have a simple DataFrame, which could just as easily be a LazyFrame...

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


I want to use a namespace to do "stuff" to both of these. For simplicity, let's just do a simple group_by sum. I could write namespaces for each separately:

@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()


But then I have to maintain code in two functions. Instead, I could convert the DataFrame to a LazyFrame in the DataFrame namespace and then call the LazyFrame namespace:

@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() #<-- convert to Lazy, do stuff, then collect
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()) #<-- only maintain code in one namespace
return ldf


a2 = a.stuff.do_stuff()
b2 = b.stuff.do_stuff()


Is this an appropriate approach, or is there a better way to deal with this?

Continue reading...