如何在Rust Polars中可靠地垂直拼接LazyFrames？

环境依赖 #

Cargo.toml:

[dependencies]
polars = { version = "0.27.2", features = ["lazy"] }

问题场景与报错 #

原本期望任意两个LazyFrames只要共同列类型兼容，缺失列自动补null（类似pandas行为）即可垂直拼接，但Polars要求二者列完全一致：

测试代码：

use polars::lazy::dsl::*;
use polars::prelude::{concat, df, DataType, IntoLazy, NamedFrom, NULL};
fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 故意将"y"放在"x"之前
    let df1 = df!["y" => &[1, 5, 17], "x" => &[1, 2, 3]].unwrap().lazy();
    let df2 = df!["x" => &[4, 5]].unwrap().lazy();
    println!(
        "{:?}",
        concat(&[df1, df2], true, true).unwrap().collect()?
    );

    Ok(())
}

报错信息：

Error: ShapeMisMatch(Owned("Could not vertically stack DataFrame. The DataFrames appended width 2 differs from the parent DataFrames width 1"))

尝试补全列后仍失败 #

给df2添加缺失的"y"列：

// 其余代码与上述示例一致
let df2 = df!["x" => &[4, 5]]
    .unwrap()
    .lazy()
    .with_column(lit(NULL).cast(DataType::Int32).alias("y"));

此时两者列和类型完全一致，仅顺序不同：

shape: (3, 2)
┌─────┬─────┐
│ y   ┆ x   │
│ --- ┆ --- │
│ i32 ┆ i32 │
╞═════╪═════╡
│ 1   ┆ 1   │
│ 5   ┆ 2   │
│ 17  ┆ 3   │
└─────┴─────┘

shape: (2, 2)
┌─────┬──────┐
│ x   ┆ y    │
│ --- ┆ ---  │
│ i32 ┆ i32  │
╞═════╪══════╡
│ 4   ┆ null │
│ 5   ┆ null │
└─────┴──────┘

但拼接仍失败，报错：

Error: SchemaMisMatch(Owned("cannot vstack: because column names in the two DataFrames do not match for left.name='y' != right.name='x'"))

显然concat()要求底层DataFrames列顺序完全一致，但LazyFrame本不应强制列顺序，因此想知道：垂直拼接这类LazyFrames的最佳方法是什么？

注：不想通过.collect()转DataFrame再堆叠，也不想手动调整列顺序。

源码补充说明 #

查看源码发现该功能暂未实现，拼接最终调用DataFrame::vstack_mut，该方法不支持缺失列或列顺序不同的情况：

pub fn vstack_mut(&mut self, other: &DataFrame) -> PolarsResult<&mut Self> {
    if self.width() != other.width() {
        if self.width() == 0 {
            self.columns = other.columns.clone();
            return Ok(self);
        }

        return Err(PolarsError::ShapeMisMatch(
            format!("Could not vertically stack DataFrame. The DataFrames appended width {} differs from the parent DataFrames width {}", self.width(), other.width()).into()
        ));
    }

    self.columns
        .iter_mut()
        .zip(other.columns.iter())
        .try_for_each::<_, PolarsResult<_>>(|(left, right)| {
            can_extend(left, right)?;
            left.append(right).expect("should not fail");
            Ok(())
        })?;
    Ok(self)
}

方法1：纯Lazy API统一列顺序 #

先获取两个LazyFrame的所有列，合并为有序列列表，再通过select()统一列顺序后拼接：

use polars::lazy::dsl::*;
use polars::prelude::{concat, df, DataType, IntoLazy, NamedFrom, NULL};
use std::collections::BTreeSet;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let df1 = df!["y" => &[1, 5, 17], "x" => &[1, 2, 3]].unwrap().lazy();
    let mut df2 = df!["x" => &[4, 5]]
        .unwrap()
        .lazy()
        .with_column(lit(NULL).cast(DataType::Int32).alias("y"));

    // 合并列名并排序，得到统一列顺序
    let schema1 = df1.schema()?;
    let schema2 = df2.schema()?;
    let all_columns: BTreeSet<_> = schema1.names().iter().chain(schema2.names().iter()).cloned().collect();
    let ordered_columns: Vec<_> = all_columns.into_iter().collect();

    // 对齐列顺序
    let df1_aligned = df1.select(&ordered_columns);
    let df2_aligned = df2.select(&ordered_columns);

    // 执行拼接
    let result = concat(&[df1_aligned, df2_aligned], true, true)?.collect()?;
    println!("{}", result);

    Ok(())
}

方法2：封装工具函数自动补全+对齐列 #

如果需要处理多个LazyFrame，可封装工具函数自动补全缺失列（按类型填充null）并统一列顺序：

use polars::lazy::dsl::*;
use polars::prelude::{LazyFrame, PolarsResult, Schema, NULL};
use std::collections::BTreeMap;

fn align_lazy_frames(frames: &[LazyFrame]) -> PolarsResult<Vec<LazyFrame>> {
    // 收集所有列的完整Schema
    let mut full_schema = BTreeMap::new();
    for frame in frames {
        let schema = frame.schema()?;
        for (name, dtype) in schema.iter() {
            full_schema.entry(name.clone()).or_insert_with(|| dtype.clone());
        }
    }

    // 对每个LazyFrame补全缺失列并统一顺序
    frames.iter().map(|frame| {
        let schema = frame.schema()?;
        let mut exprs = Vec::new();

        for (name, dtype) in &full_schema {
            exprs.push(
                if schema.contains(name) {
                    col(name)
                } else {
                    lit(NULL).cast(dtype.clone()).alias(name)
                }
            );
        }

        Ok(frame.select(exprs))
    }).collect()
}

// 使用示例
fn main() -> Result<(), Box<dyn std::error::Error>> {
    let df1 = df!["y" => &[1, 5, 17], "x" => &[1, 2, 3]].unwrap().lazy();
    let df2 = df!["x" => &[4, 5]].unwrap().lazy();

    let aligned_frames = align_lazy_frames(&[df1, df2])?;
    let result = concat(&aligned_frames, true, true)?.collect()?;
    println!("{}", result);

    Ok(())
}

方法3：Polars 0.30+版本直接用官方新函数 #

Polars 0.30及以上版本新增了concat_with_nulls，可直接实现自动补全缺失列、忽略列顺序的垂直拼接：

// Polars 0.30+ 可用
use polars::prelude::{concat_with_nulls, df, IntoLazy};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let df1 = df!["y" => &[1, 5, 17], "x" => &[1, 2, 3]].unwrap().lazy();
    let df2 = df!["x" => &[4, 5]].unwrap().lazy();

    let result = concat_with_nulls(&[df1, df2])?.collect()?;
    println!("{}", result);

    Ok(())
}

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