Polars LazyFrame通过PartitionByKey写入S3性能异常缓慢的原因排查

我想搞清楚为什么用PartitionByKey把LazyFrame写入S3时，性能比其他方式差这么多。以下是我用来对比本地磁盘和S3写入性能的测试脚本：

import time

import polars as pl
import numpy as np

l = 600_000_000

data = pl.DataFrame({"a": np.random.randint(0, 100, l),
    "b": np.random.randint(0, 100, l)})

t0 = time.time()
data.write_parquet("/tmp/in-memory.parquet",
    compression="zstd", compression_level=6)
t1 = time.time()

print(f"in-memory, local write: {t1-t0:0.2f}s")

t0 = time.time()
data.write_parquet("s3://bucket/in-memory.parquet",
    compression="zstd", compression_level=6)
t1 = time.time()

print(f"in-memory, s3 write: {t1-t0:0.2f}s")

lazy = data.lazy()

t0 = time.time()
lazy.sink_parquet(
        pl.PartitionByKey(
            base_path="/tmp/lazy",
            file_path=lambda ctx: f"{ctx.keys[0]}.parquet",
            by="a",
            include_key=True),
        mkdir=True,
        compression="zstd",
        compression_level=6,
        maintain_order=False,
        )
t1 = time.time()

print(f"lazy, local write: {t1-t0:0.2f}s")

t0 = time.time()
lazy.sink_parquet(
        pl.PartitionByKey(
            base_path="s3://bucket/lazy/",
            file_path=lambda ctx: f"{ctx.keys[0]}.parquet",
            by="a",
            include_key=True),
        mkdir=True,
        compression="zstd",
        compression_level=6,
        maintain_order=False,
        )
t1 = time.time()

print(f"lazy, s3 write: {t1-t0:0.2f}s")

我在AWS EC2 r8g.xlarge实例上运行该脚本，实例配置：4核CPU，gp3磁盘（默认125MB/s带宽），网络带宽最高12.5Gbps（比磁盘速度快），使用Polars版本1.34.0。测试结果如下：

in-memory, local write: 11.13s
in-memory, s3 write: 5.34s
lazy, local write: 3.73s
lazy, s3 write: 18.62s

前三项结果符合预期：多核并行或更快的网络能提升性能。但最后一项——LazyFrame通过PartitionByKey写入S3的耗时却异常长。按我的理解，这应该是最快的方式才对，因为它既能利用所有核心，又能借助比磁盘更快的网络。我是不是操作有误？或者遗漏了什么配置？

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