减少 unique() 中的 Polars 内存消耗

Question

我有一个适合 RAM 的数据集，但当我运行某些方法（例如df.unique(). 我的笔记本电脑有 16GB 内存。我正在使用 14GB RAM 运行 WSL。我使用的是 Polars 版本 0.18.4。运行df.estimated_size()表明，当我读入数据集时，我的数据集约为 6GB。我的数据的架构是

index: Int64
first_name: Utf8
last_name: Utf8
race: Utf8
pct_1: Float64
pct_2: Float64
pct_3: Float64
pct_4: Float64

index: Int64
first_name: Utf8
last_name: Utf8
race: Utf8
pct_1: Float64
pct_2: Float64
pct_3: Float64
pct_4: Float64

但是，我无法在不收到 SIGKILL 信号的情况下执行诸如.unique()、.drop_nulls()等任务。我正在使用 LazyFrames。

例如，

size = pl.read_parquet("data.parquet").estimated_size()
df = pl.scan_parquet("data.parquet") # use LazyFrames

导致内存不足错误。我可以通过编写自定义函数来回避这个问题。

df = df.drop_nulls().collect(streaming=True)

我很好奇为什么后者有效而前者无效，因为数据集的最大版本（当我最初读入它时）适合 RAM。

不幸的是，我无法想出类似的技巧来做同样的事情.unique()。我可以做些什么来.unique()减少内存占用吗？我努力了：

def iterative_drop_nulls(expr: pl.Expr, subset: list[str]) -> pl.LazyFrame:
    for col in subset:
        expr = expr.filter(~pl.col(col).is_null())

    return expr

df = df.pipe(iterative_drop_nulls, ["col1", "col2"]).collect()

和

df = df.lazy().unique(cols).collect(streaming=True)

编辑：

我希望有一个更好的答案，但现在我正在使用

def unique(df: pl.DataFrame, subset: list[str], n_rows: int = 100_000) -> pl.DataFrame:
    parts = []
    for slice in df.iter_slices(n_rows=n_rows):
        parts.append(df.unique(slice, subset=subset))

    return pl.concat(parts)

总的来说，我发现 Polars 的内存效率比 Pandas 更高，但也许这是 Polars 可以改进的一个领域？奇怪的是，如果我使用

df = pl.from_pandas(
    df.collect()
    .to_pandas()
    .drop_duplicates(subset=["col1", "col2"])
)

我也遇到了同样的内存错误，所以也许这是 Pyarrow 的问题。