如何在Rust Polars中可靠地连接LazyFrame

col17t5w  于 2023-03-08  发布在  其他
关注(0)|答案(1)|浏览(248)

Cargo.toml:

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

我希望任何两个LazyFrame都可以垂直连接,只要它们共有的列具有相同或可提升的数据类型,缺失的列作为null添加(就像Pandas做的那样)。但显然它们需要具有相同的列:

use polars::lazy::dsl::*;
use polars::prelude::{concat, df, DataType, IntoLazy, NamedFrom, NULL};
fn main() -> Result<(), Box<dyn std::error::Error>> {
    // "y" intentionally comes before "x" here
    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"))有关。
我尝试将缺少的"y"列添加到df2

// everything but this line is the same as above
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()要求列在底层DataFrame中的顺序相同。
但是我认为在LazyFrames中不可能强制任何特定的列顺序(实际上也不需要强制,因为列顺序应该是无关紧要的)。那么,垂直连接这两个LazyFrames的最佳方法是什么呢?
如果可能的话,我不希望将它们分别.collect()到 Dataframe 中,然后vstack Dataframe 并在结果上调用.lazy();这似乎是不必要的复杂。如果我做了.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)
}
l7wslrjt

l7wslrjt1#

一旦你知道去哪里找,答案就变得相当简单了,使用diagonal_concat功能,你可以解锁diag_concat_lf(以及渴望的diag_concat_df):

pub fn diag_concat_lf<L>(lfs: L, rechunk: bool, parallel: bool) -> PolarsResult<LazyFrame>
where
    L: AsRef<[LazyFrame]>,
{ ... }

pub fn diag_concat_df(dfs: &[DataFrame]) -> PolarsResult<dataframe> { ... }

相关问题