使用 PyArrow 将 DataFrame 写入 Polars 中的分区镶木地板对象

Question

我的目录中有 12 个镶木地板文件，其中包含我尝试使用 Polars 和 PyArrow 写入分区对象的匹配列。我正在迭代目录中的每个文件并将其作为 LazyFrame 读取。然后，我迭代 DataFrame 列表并将它们写入分区对象。每个 DataFrame 的估计大小约为 1GB，所有 DataFrame 连接起来约为 10GB。该过程使用约 15GB RAM，在一小时内完成。

我尝试使用以下代码来做到这一点：

all_lazyframes: list[pl.LazyFrame] = []

for file in glob.glob(input_path):
    lazyframe: pl.LazyFrame = pl.scan_parquet(file)
    all_lazyframes.append(lazyframe)

dataframes: list[pl.DataFrame] = pl.collect_all(all_lazyframes)

for output in dataframes:
    output.write_parquet(
        output_path,
        use_pyarrow=True,
        pyarrow_options={"partition_cols": ["part"]},
    )

生成的分区对象具有以下结构：

partitioned_object/
  part=a/
       data0.parquet
       data1.parquet
       ...
  part=b/
       data0.parquet
       data1.parquet
    ...

该对象的大小约为 250GB。我的问题是，当输入数据总共只有 ~10GB 时，为什么分区对象这么大？有没有更有效的方法来实现这一目标？

Answer 1

我通过在 ds.write_dataset() 函数中指定每组的行大小解决了这个问题。目前，极坐标方法需要更长的时间，并且在使用 PyArrow 选项时无法指定每组的行数。

ds.write_dataset(
        data,
        output_path,
        format="parquet",
        min_rows_per_group=1000000,
        max_rows_per_group=1000000,
        partitioning=ds.partitioning(pa.schema([("type", pa.string())])),
        existing_data_behavior="overwrite_or_ignore"
    )

Polars 方法需要更长的时间，并且使用 PyArrow 选项时 row_group_size 选项不起作用：

    output.write_parquet(
       file=output_path,
       use_pyarrow=True,
       pyarrow_options={"partition_cols": partition_cols},
       row_group_size=1000000,
    )