如何扩展Polars API使其同时支持DataFrame与LazyFrame

你的第二种方法（DataFrame转LazyFrame复用命名空间逻辑）是完全可行的，而且能有效减少代码重复、降低维护成本，是很实用的方案。

为什么这种方法合适？ #

• Polars中DataFrame.lazy()是轻量转换，不会立即计算，性能开销可以忽略
• 只需要维护LazyFrame版本的命名空间逻辑，所有DataFrame的调用都复用这部分代码，避免了两份逻辑不一致的风险
• 最终通过.collect()转回DataFrame，完全符合Polars的延迟执行设计理念

更优的实现方式：共享核心逻辑 #

除了转LazyFrame的方式，还可以通过提取核心逻辑或使用基类来实现代码复用，让两种命名空间类共享同一套操作逻辑：

方法1：提取通用核心函数

把操作逻辑抽成独立函数，让两个命名空间类都调用它：

import polars as pl

def _do_stuff(frame):
    # 核心逻辑，同时兼容DataFrame和LazyFrame
    return frame.group_by("col1").agg(pl.col("*").sum())

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

    def do_stuff(self) -> pl.DataFrame:
        return _do_stuff(self._df)

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

    def do_stuff(self) -> pl.LazyFrame:
        return _do_stuff(self._ldf)

# 测试
a = pl.DataFrame({"col1": ["A", "B", "A", "A", "B"], "col2": [1, 2, 3, 4, 5]})
b = a.lazy()
a1 = a.stuff.do_stuff()
b1 = b.stuff.do_stuff()

方法2：使用基类继承

通过基类封装通用方法，两种命名空间类继承基类并初始化各自的框架对象：

import polars as pl

class BaseStuffNamespace:
    def do_stuff(self):
        # 通用操作逻辑
        return self._frame.group_by("col1").agg(pl.col("*").sum())

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

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

# 测试
a = pl.DataFrame({"col1": ["A", "B", "A", "A", "B"], "col2": [1, 2, 3, 4, 5]})
b = a.lazy()
a1 = a.stuff.do_stuff()
b1 = b.stuff.do_stuff()

方法选择建议 #

• 如果你的操作逻辑更适合在延迟执行模式下优化（比如复杂多步骤查询），优先选择你原来的转LazyFrame方案
• 如果操作逻辑对DataFrame和LazyFrame完全通用，使用核心函数或基类继承的方式更直观，避免额外的转换步骤

内容的提问来源于stack exchange，提问作者blaylockbk