Function window_using_aggregates

datafusion/core/tests/dataframe/mod.rs:1150–1233

()

#[tokio::test]
async fn window_using_aggregates() -> Result<()> {
    // build plan using DataFrame API
    let df = test_table().await?.filter(col("c1").eq(lit("a")))?;
    let mut aggr_expr = vec![
        (
            datafusion_functions_aggregate::first_last::first_value_udaf(),
            "first_value",
        ),
        (
            datafusion_functions_aggregate::first_last::last_value_udaf(),
            "last_val",
        ),
        (
            datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf(),
            "approx_distinct",
        ),
        (
            datafusion_functions_aggregate::approx_median::approx_median_udaf(),
            "approx_median",
        ),
        (
            datafusion_functions_aggregate::median::median_udaf(),
            "median",
        ),
        (datafusion_functions_aggregate::min_max::max_udaf(), "max"),
        (datafusion_functions_aggregate::min_max::min_udaf(), "min"),
    ]
    .into_iter()
    .map(|(func, name)| {
        let w = WindowFunction::new(
            WindowFunctionDefinition::AggregateUDF(func),
            vec![col("c3")],
        );

        Expr::from(w)
            .null_treatment(NullTreatment::IgnoreNulls)
            .order_by(vec![col("c2").sort(true, true), col("c3").sort(true, true)])
            .window_frame(WindowFrame::new_bounds(
                WindowFrameUnits::Rows,
                WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
                WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))),
            ))
            .build()
            .unwrap()
            .alias(name)
    })
    .collect::<Vec<_>>();
    aggr_expr.extend_from_slice(&[col("c2"), col("c3")]);

    let df: Vec<RecordBatch> = df.select(aggr_expr)?.collect().await?;

    assert_snapshot!(
        batches_to_sort_string(&df),
        @r"
    +-------------+----------+-----------------+---------------+--------+-----+------+----+------+
    | first_value | last_val | approx_distinct | approx_median | median | max | min  | c2 | c3   |
    +-------------+----------+-----------------+---------------+--------+-----+------+----+------+
    |             |          |                 |               |        |     |      | 1  | -85  |