Python上下文管理器监控Polars DataFrame行变化失效问题

问题根源 #

Polars DataFrame是不可变对象，调用unique()、filter()这类转换方法时，会返回一个全新的DataFrame实例，而非修改原对象。

你的代码中，上下文管理器接收的是原DataFrame的引用，但在with块里执行df = df.unique(...)只是将局部变量df重新指向了新的DataFrame对象，上下文管理器内部的df仍然指向最初的旧对象，因此finally块统计的始终是原DataFrame的行数（3行），导致计数结果不符合预期。

解决方案 #

方案1：用可变容器包装DataFrame #

利用列表这类可变对象来传递DataFrame，这样在with块内修改容器内的元素时，上下文管理器能感知到新的DataFrame实例：

import polars as pl
from contextlib import contextmanager

def count_rows(df: pl.DataFrame) -> int:
    """统计Polars DataFrame的行数"""
    return df.select(pl.count()).item()

@contextmanager
def log_row_count_change(df_container: list, action_desc: str = '', df_name: str = 'df') -> None:
    try:
        df_before = df_container[0]
        row_count_before = count_rows(df_before)
        print(f"执行'{action_desc}'前，'{df_name}'的行数：{row_count_before:,}")
        yield
    finally:
        df_after = df_container[0]
        row_count_after = count_rows(df_after)
        row_count_change = row_count_after - row_count_before
        row_count_change_pct = row_count_change / row_count_before * 100
        print(f"执行'{action_desc}'后，'{df_name}'的行数变化：{row_count_change:,}行（{row_count_before:,} → {row_count_after:,}），变化率：{row_count_change_pct:.2f}%")

# 测试代码
df = pl.DataFrame({"a":[1,1,2], "b":[2,2,3], "c":[1,2,3]})
df_container = [df]

with log_row_count_change(df_container, '按列a去重', 'df'):
    df_container[0] = df_container[0].unique(subset=['a'])

方案2：让上下文管理器接收操作逻辑 #

调整上下文管理器，使其返回一个可执行操作的函数，直接在内部处理新DataFrame的计数：

import polars as pl
from contextlib import contextmanager

def count_rows(df: pl.DataFrame) -> int:
    return df.select(pl.count()).item()

@contextmanager
def log_row_count_change(df: pl.DataFrame, action_desc: str = '', df_name: str = 'df'):
    row_count_before = count_rows(df)
    print(f"执行'{action_desc}'前，'{df_name}'的行数：{row_count_before:,}")
    
    def apply_action(operation):
        new_df = operation(df)
        row_count_after = count_rows(new_df)
        row_count_change = row_count_after - row_count_before
        row_count_change_pct = row_count_change / row_count_before * 100
        print(f"执行'{action_desc}'后，'{df_name}'的行数变化：{row_count_change:,}行（{row_count_before:,} → {row_count_after:,}），变化率：{row_count_change_pct:.2f}%")
        return new_df
    
    yield apply_action

# 测试代码
df = pl.DataFrame({"a":[1,1,2], "b":[2,2,3], "c":[1,2,3]})

with log_row_count_change(df, '按列a去重', 'df') as apply:
    df = apply(lambda d: d.unique(subset=['a']))

运行结果 #

两种方案都会输出符合预期的结果：

执行'按列a去重'前，'df'的行数：3
执行'按列a去重'后，'df'的行数变化：-1行（3 → 2），变化率：-33.33%

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