求Rust Polars中对应Python LazyFrame.take_every的实现及降采样方案

Rust Polars 目前没有直接对应 Python LazyFrame.take_every() 的API，但可以通过行索引过滤的方式实现相同逻辑，同时结合获取LazyFrame长度来计算合适的采样步长，确保最终行数少于n行。

实现步骤与代码 #

1. 获取总行数：通过惰性计算的pl.len()获取总行数，再执行collect()拿到具体数值（LazyFrame是惰性的，必须触发计算才能得到实际长度）。
2. 计算采样步长：为保证采样后行数严格小于n，步长用向上取整公式计算：step = (total_rows + n - 2) // (n - 1)。如果允许最终行数等于n，可改用step = (total_rows + n - 1) // n。
3. 过滤行：用int_range生成行索引，通过取模运算筛选每step行保留一行。

use polars::prelude::*;

fn downsample_lazyframe(lf: LazyFrame, max_rows: usize) -> PolarsResult<LazyFrame> {
    // 获取LazyFrame总行数
    let total_rows = lf.clone()
        .select(pl.len())
        .collect()?
        .column("len")?
        .i64()?
        .get(0)
        .unwrap_or(0) as usize;

    // 如果总行数已经符合要求，直接返回原数据
    if total_rows < max_rows {
        return Ok(lf);
    }

    // 计算步长，确保采样后行数 < max_rows
    let step = (total_rows + max_rows - 2) // (max_rows - 1);

    // 模拟take_every的过滤逻辑
    Ok(lf.filter(
        pl.int_range(0, pl.len())
            .mod(lit(step as i64))
            .eq(lit(0))
    ))
}

// 示例调用
fn main() -> PolarsResult<()> {
    // 创建测试用LazyFrame
    let test_df = df!(
        "data" => (0..105).collect::<Vec<i32>>()
    )?;
    let lazy_frame = test_df.lazy();

    // 降采样到少于20行
    let downsampled = downsample_lazyframe(lazy_frame, 20)?;
    let result_df = downsampled.collect()?;

    println!("降采样后行数: {}", result_df.height()); // 输出18（105/(20-1)=5.526，步长6，105/6=17.5 → 18行，小于20）
    println!("{}", result_df);

    Ok(())
}

关键细节 #

• 克隆原LazyFrame计算长度，避免破坏原有的惰性计算链。
• 步长计算的向上取整逻辑确保最终行数严格小于设定的n值。
• int_range生成的行索引是惰性计算的，不会提前加载全量数据，符合LazyFrame的设计理念。

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