如何快速从Polars LazyFrame（超大CSV）中随机选取单条数据？

问题根源 #

CSV是行式文本文件，没有内置的行位置索引。使用slice(index, 1)时，Polars需要从文件开头逐行扫描到目标索引位置，当索引接近文件末尾时，扫描量极大，导致速度骤降。

最优方案：预生成行偏移索引 #

预先遍历一次CSV文件，记录每一行的起始字节偏移量并保存为小索引文件。之后读取随机行时，直接通过偏移量跳转到目标位置，无需扫描前置内容。

步骤1：生成行偏移索引 #

import csv
import polars as pl

def build_row_offsets(csv_path, offsets_path):
    offsets = []
    with open(csv_path, 'rb') as f:
        # 匹配has_header=False设置，第一行即为数据，记录初始偏移
        offsets.append(f.tell())
        # 用csv模块解析行，避免跨行字段、换行符格式问题
        reader = csv.reader((line.decode('utf-8') for line in f))
        for _ in reader:
            offsets.append(f.tell())
    # 将偏移量保存为Parquet文件（占用空间小、读取效率高）
    pl.DataFrame({'offset': offsets}).write_parquet(offsets_path)

• 此函数只需运行一次，生成的索引文件仅占原CSV的极小比例（通常不到1%）。
• 依赖csv模块解析行，确保准确识别所有合法行结构。

步骤2：快速读取随机行 #

def get_random_row(csv_path, offsets_path):
    # 加载预生成的偏移量索引
    offsets_df = pl.read_parquet(offsets_path)
    # 随机选择一个有效行索引
    random_idx = pl.int_range(0, len(offsets_df)-1).sample(n=1)[0]
    target_offset = offsets_df['offset'][random_idx]
    # 通过字节偏移直接跳转，仅读取目标行
    row = pl.scan_csv(
        csv_path,
        has_header=False,
        skip_rows_by_bytes=target_offset,
        n_rows=1
    ).collect().row(0)
    return row

• skip_rows_by_bytes让Polars直接跳转到指定字节位置，完全跳过前置内容的扫描。
• n_rows=1限制仅读取目标行，避免多余IO操作。

备选方案：水库采样（适合无指定索引的随机采样） #

如果只需随机获取一条数据，无需指定具体行索引，可以使用Polars的水库采样策略：

random_row = df.sample(n=1, strategy="reservoir").collect().row(0)

• 水库采样支持流式处理，无需加载全量数据，但第一次采样仍需遍历整个文件，适合需要多次重复随机采样的场景。

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