Function with_column_window_functions

datafusion/core/tests/dataframe/mod.rs:168–203

()

#[tokio::test]
async fn with_column_window_functions() -> DataFusionResult<()> {
    let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);

    let batch = RecordBatch::try_new(
        Arc::new(schema.clone()),
        vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
    )?;

    let ctx = SessionContext::new();

    let provider = MemTable::try_new(Arc::new(schema), vec![vec![batch]])?;
    ctx.register_table("t", Arc::new(provider))?;

    // Define test cases: (expr builder, alias name)
    let test_cases: Vec<WindowFnCase> = vec![
        (|| lag(col("a"), Some(1), None), "lag_val"),
        (|| lead(col("a"), Some(1), None), "lead_val"),
        (row_number, "row_num"),
    ];

    for (make_expr, alias) in test_cases {
        let df = ctx.table("t").await?;
        let expr = make_expr();
        let df_with = df.with_column(alias, expr)?;
        let df_schema = df_with.schema().clone();

        assert!(
            df_schema.has_column_with_unqualified_name(alias),
            "Schema does not contain expected column {alias}",
        );

        assert_eq!(2, df_schema.columns().len());
    }

    Ok(())
}

#[tokio::test]
async fn test_coalesce_schema() -> Result<()> {