polars-lazy 0.55.2

Lazy query engine for the Polars DataFrame library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

use super::*;

#[test]
fn test_schema_update_after_projection_pd() -> PolarsResult<()> {
    let df = df![
        "a" => [1],
        "b" => [1],
        "c" => [1],
    ]?;

    let q = df
        .lazy()
        .with_column(col("a").implode(true))
        .explode(
            by_name(["a"], true, false),
            ExplodeOptions {
                empty_as_null: true,
                keep_nulls: true,
            },
        )
        .select([cols(["a", "b"]).as_expr()]);

    // run optimizations
    // Get the explode node
    let IRPlan {
        lp_top,
        lp_arena,
        expr_arena: _,
    } = q.to_alp_optimized()?;

    // assert the schema has been corrected with the projection pushdown run
    let lp = lp_arena.get(lp_top);
    assert!(matches!(
        lp,
        IR::MapFunction {
            function: FunctionIR::Explode { .. },
            ..
        }
    ));

    let schema = lp.schema(&lp_arena).into_owned();
    let mut expected = Schema::default();
    expected.with_column("a".into(), DataType::Int32);
    expected.with_column("b".into(), DataType::Int32);
    assert_eq!(schema.as_ref(), &expected);

    Ok(())
}