Polars流处理计算大数据集外积的内存问题及优化咨询

针对你用Polars处理1400万行数据集时遇到的内存问题，以下是具体解答：

1. 代码优化建议 #

• 避免重复调用collect_schema()：原代码多次调用df.collect_schema().names()会重复读取文件头部获取schema，建议提前一次性获取并复用：

  import polars as pl
  
  # 提前获取原始schema并转小写
  raw_schema = pl.scan_csv("dummy_file.csv").collect_schema()
  lower_cols = {col: col.lower() for col in raw_schema.names()}
  
  # 初始化惰性DataFrame
  df = pl.scan_csv("dummy_file.csv", schema=raw_schema).rename(lower_cols)
  

• 直接构造特征列表达式，减少中间列：不要先添加72个sin/cos列再拼接成列表，直接在concat_list中生成所需特征，避免中间列占用内存：

  # 构造包含原始、sin、cos的特征表达式列表
  feature_exprs = (
      [pl.col(f) for f in features] +
      [pl.col(f).sin() for f in features] +
      [pl.col(f).cos() for f in features]
  )
  
  # 直接生成features列表列，无需留存中间sin/cos列
  df = df.select(['x','y','z','t'] + features).with_columns(
      pl.concat_list(feature_exprs).alias("features")
  )
  

• 限制流处理批次内存：通过环境变量设置单批次最大内存阈值，强制Polars使用更小的批次处理：

  import os
  # 设置单批次最大内存为2GB（根据16GB内存调整）
  os.environ["POLARS_MAX_MEMORY_BYTES"] = str(2 * 1024**3)
  

• 使用紧凑数据类型：若精度允许，将浮点型从float64转为float32，直接减半内存占用：

  df = df.with_columns([pl.col(f).cast(pl.Float32) for f in features])
  

2. 流处理未生效的原因 #

• 单个批次内存过载：外积计算每行生成3*36 * 3*36 = 11664个元素，若默认批次过大（如10万行），单批次内存占用可达~9.3GB，直接超出16GB内存上限。
• 中间列内存留存：原代码先添加72个sin/cos列再拼接列表，这些中间列会在每个批次中占用额外内存，直到后续步骤被丢弃。
• 嵌套List的内存开销：Polars对大量嵌套List的拼接操作会产生额外内存开销，临时内存无法及时释放导致占用持续增长。

3. 内存高效的外积计算替代方案 #

• 用list.eval简化外积计算：通过行内List求值避免生成多个中间List，直接计算外积并展平：

  df = df.with_columns(
      pl.col("features").list.eval(
          pl.element() * pl.col("features"),
          return_dtype=pl.List(pl.Float32)
      ).list.flatten().alias("outer_products")
  )
  

  该方式对features列表的每个元素与整个列表相乘生成子列表，再展平为一维List，相比循环list.get(i)减少了中间List的创建开销。
• 手动分块处理：若流处理仍无法解决，可将数据集拆分为小文件逐个处理，写完后显式释放内存：

  import gc
  
  # 假设已将CSV拆分为多个小文件
  for chunk_file in chunk_files:
      df_chunk = pl.read_csv(chunk_file).rename(lower_cols)
      # 执行特征转换与外积计算
      df_chunk = df_chunk.select(['x','y','z','t'] + features).with_columns(
          pl.concat_list(feature_exprs).alias("features")
      ).with_columns(
          pl.col("features").list.eval(pl.element() * pl.col("features")).list.flatten().alias("outer_products")
      )
      df_chunk.write_parquet('test.parquet', append=True)
      del df_chunk
      gc.collect()
  

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