With polars dataframe, is it possible to serialize/deserialize only the transformations and not the original dataframe?

user_1538095 0
Viewed 224

I'd like to be able to serialize the transformations, like group by or so, and then apply them to an existing dataframe. I'd rather not have to build the LogicalPlan recursively, as this would mean having to manually implement and keep up to date each transformation types.

rust-polars rust
1 Answers

While you can indeed serialize and deserialize a logical plan, I don't see how to apply a logical plan to an existing dataframe — a method like df.apply(logical_plan) doesn't seem to exist.

First, Cargo.toml:

[dependencies]
    // There's a serde-related bug in polars-lazy@0.22.7 that's fixed on master
    // but hasn't been released yet
    polars = { git = "https://github.com/pola-rs/polars", rev = "25b58e2ec82d95e4ec6120fa4f5c49e7d1083f39", features = [
        "lazy",
        "serde",
        "serde-lazy",
    ] }
    serde_json = "1.0.82"

Here's what I have:

fn main() {
    use polars::prelude::*;

    // Create (eager) DataFrame
    let df = DataFrame::new(vec![
        Series::new("Fruit", &["Apple", "Apple", "Pear"]),
        Series::new("Color", &["Red", "Green", "Green"]),
        Series::new("Quantity", &[1, 2, 3]),
    ])
    .unwrap();

    // Create LazyFrame
    let lazy_df = df.lazy().filter(col("Quantity").gt(1)).select(&[
        col("Fruit").last(),
        col("Color").n_unique(),
        col("Quantity").sum(),
    ]);

    // Serialize logical plan
    serde_json::to_writer(
        std::fs::File::create("logical_plan.json").expect("could not create file"),
        &lazy_df.logical_plan,
    )
    .expect("could not serialize");

    // Deserialize logical plan
    let logical_plan: LogicalPlan = {
        serde_json::from_reader(
            std::fs::File::open("logical_plan.json").expect("could not open file"),
        )
        .expect("could not deserialize")
    };

    println!("LazyFrame:\n{:?}", lazy_df.clone().collect().unwrap());
    println!("Its logical plan:\n{:?}", lazy_df.logical_plan);
    println!("Deserialized plan:\n{:?}", logical_plan);
}

As expected, this prints the following:

LazyFrame:
shape: (1, 3)
┌───────┬───────┬──────────┐
│ Fruit ┆ Color ┆ Quantity │
│ ---   ┆ ---   ┆ ---      │
│ str   ┆ u32   ┆ i32      │
╞═══════╪═══════╪══════════╡
│ Pear  ┆ 1     ┆ 5        │
└───────┴───────┴──────────┘
Its logical plan:
SELECT 3 COLUMNS: [col("Fruit").last(), col("Color").n_unique(), col("Quantity").sum()]
FROM
FILTER [(col("Quantity")) > (1i32)]
FROM
DATAFRAME(in-memory): ["Fruit", "Color", "Quantity"];
        project */3 columns     |       details: None;
        selection: "None"


Deserialized plan:
SELECT 3 COLUMNS: [col("Fruit").last(), col("Color").n_unique(), col("Quantity").sum()]
FROM
FILTER [(col("Quantity")) > (1i32)]
FROM
DATAFRAME(in-memory): ["Fruit", "Color", "Quantity"];
        project */3 columns     |       details: None;
        selection: "None"

So all that's left to do is to apply this logical plan to an existing frame, but like I said I cannot find such a function. If it truly doesn't exist then it may be worth opening an issue on polars's GitHub and seeing if they'll add this functionality.

user_54964330
Related