使用Polars从S3读取Parquet文件时内存占用激增问题咨询

当前内存占用是否正常？ #

这种内存占用完全不正常。600万行仅含短文本的Parquet数据，即使是未压缩的内存表示，正常占用也应该在几百MB级别，10GB的内存消耗明显远超合理范围，说明读取过程中存在可优化的环节。

可能的异常原因 #

• PyArrow Dataset默认的类型推断开销：PyArrow在自动推断列类型时，可能会加载更多数据样本，哪怕只有短文本，若存在隐含的类型歧义也会导致额外内存占用。
• S3文件系统的缓存配置：PyArrow的S3FileSystem默认启用较大缓存空间，若未限制，可能会缓存大量文件内容到内存。
• Polars与PyArrow Dataset的协同问题：scan_pyarrow_dataset的默认配置可能未充分利用Polars的懒加载优化，导致提前加载不必要的数据。

优化方案 #

• 显式指定数据Schema：提前定义列的具体类型，避免PyArrow自动推断时的额外内存开销。示例代码：

  import pyarrow as pa
  import polars as pl
  from pyarrow.dataset import dataset as ds
  from pyarrow.fs import S3FileSystem
  
  # 按实际列定义schema
  custom_schema = pa.schema([
      ("column_1", pa.string()),
      ("column_2", pa.int32()),
      # 补充其他列的类型定义
  ])
  
  s3 = S3FileSystem()
  dataset = ds(f"my_bucket/myfiles/", filesystem=s3, schema=custom_schema)
  df = pl.scan_pyarrow_dataset(dataset).collect()
  

• 改用Polars原生Parquet扫描：Polars对Parquet的读取有专门优化，直接使用scan_parquet替代PyArrow Dataset，往往能降低内存占用：

  df = pl.scan_parquet("s3://my_bucket/myfiles/*.parquet").collect()
  

• 限制S3缓存大小：初始化S3FileSystem时设置缓存上限，避免无限制缓存文件内容：

  s3 = S3FileSystem(cache_options={"capacity": 100 * 1024 * 1024})  # 限制缓存为100MB
  

• 调整批量读取参数：在扫描时指定batch_size，让数据分批加载，降低内存峰值：

  df = pl.scan_pyarrow_dataset(dataset, batch_size=1_000_000).collect()
  

• 关闭不必要的缓存：禁用Polars的字符串缓存（如果不需要），减少内存占用：

  pl.enable_string_cache(False)
  

读取更大文件的可行性 #

如果不做优化，直接读取600-700MB的Parquet文件，大概率会出现内存溢出问题。但通过上述优化手段，把内存峰值控制在2GB以内是完全可行的。建议先优化当前的4个文件读取逻辑，验证内存占用降至合理范围后，再处理更大的文件。

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