enable_join_dynamic_filter_pushdown silently drops rows: the pushed filter is remapped by column NAME and lands on the wrong same-named column #25244

use std::sync::Arc;

use arrow::array::{ArrayRef, StringArray};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;
use datafusion::execution::context::SessionConfig;
use datafusion::logical_expr::{Expr, JoinType, LogicalPlan, LogicalPlanBuilder};
use datafusion::physical_optimizer::PhysicalOptimizerRule;
use datafusion::prelude::*;
use parquet::arrow::ArrowWriter;

fn write_parquet(path: &std::path::Path, schema: Arc<Schema>, batch: &RecordBatch) {
    let file = std::fs::File::create(path).unwrap();
    let mut w = ArrowWriter::try_new(file, schema, None).unwrap();
    w.write(batch).unwrap();
    w.close().unwrap();
}

// ta: single row (id = "a1", ty = "x1")  -> id range [a1,a1] does NOT contain x1
// tb: single row (id = "x1")
fn make_tables() -> (tempfile::TempDir, String, String) {
    let dir = tempfile::tempdir().unwrap();

    let sa = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Utf8, false),
        Field::new("ty", DataType::Utf8, false),
    ]));
    let b = RecordBatch::try_new(
        Arc::clone(&sa),
        vec![
            Arc::new(StringArray::from(vec!["a1"])) as ArrayRef,
            Arc::new(StringArray::from(vec!["x1"])) as ArrayRef,
        ],
    )
    .unwrap();
    let ta = dir.path().join("ta.parquet");
    write_parquet(&ta, sa, &b);

    let sb = Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)]));
    let b = RecordBatch::try_new(
        Arc::clone(&sb),
        vec![Arc::new(StringArray::from(vec!["x1"])) as ArrayRef],
    )
    .unwrap();
    let tb = dir.path().join("tb.parquet");
    write_parquet(&tb, sb, &b);

    (
        dir,
        ta.to_string_lossy().into_owned(),
        tb.to_string_lossy().into_owned(),
    )
}

async fn ctx(dynamic: bool, ta: &str, tb: &str) -> SessionContext {
    let mut config = SessionConfig::new();
    config.options_mut().execution.target_partitions = 1;
    config.options_mut().execution.collect_statistics = true;
    config.options_mut().execution.parquet.schema_force_view_types = false;
    config.options_mut().optimizer.join_reordering = false;
    config.options_mut().optimizer.enable_join_dynamic_filter_pushdown = dynamic;
    let ctx = SessionContext::new_with_config(config);
    ctx.register_parquet("ta", ta, ParquetReadOptions::default()).await.unwrap();
    ctx.register_parquet("tb", tb, ParquetReadOptions::default()).await.unwrap();
    ctx
}

fn c(rel: &str, name: &str) -> Expr {
    Expr::Column(datafusion::common::Column::new(Some(rel), name))
}
fn k(rel: &str, name: &str) -> datafusion::common::Column {
    datafusion::common::Column::new(Some(rel), name)
}

fn build_logical(scan_a: LogicalPlan, scan_b: LogicalPlan) -> LogicalPlan {
    let aliased = |scan: LogicalPlan, name: &str, exprs: Vec<Expr>| {
        LogicalPlanBuilder::from(scan)
            .alias(name).unwrap()
            .project(exprs).unwrap()
            .build().unwrap()
    };
    let a = aliased(scan_a, "a", vec![c("a", "id"), c("a", "ty")]);
    let b = aliased(scan_b.clone(), "b", vec![c("b", "id")]);
    let cc = aliased(scan_b.clone(), "c", vec![c("c", "id")]);

    // A = a JOIN b ON a.ty = b.id -> [a.id, a.ty, b.id]
    let a_join = LogicalPlanBuilder::from(a)
        .join_detailed(
            b, JoinType::Inner,
            (vec![k("a", "ty")], vec![k("b", "id")]),
            None, datafusion::common::NullEquality::NullEqualsNothing,
        ).unwrap().build().unwrap();

    // A2 = [a.id, b.id] -> TWO output columns both named `id`
    let a2 = LogicalPlanBuilder::from(a_join)
        .project(vec![c("a", "id"), c("b", "id")])
        .unwrap().build().unwrap();

    // J = A2 JOIN c ON b.id = c.id -> [a.id, b.id, c.id]
    let j = LogicalPlanBuilder::from(a2)
        .join_detailed(
            cc, JoinType::Inner,
            (vec![k("b", "id")], vec![k("c", "id")]),
            None, datafusion::common::NullEquality::NullEqualsNothing,
        ).unwrap().build().unwrap();

    // s = filter(tb, id = 'x1') -> [s.id]
    let s = LogicalPlanBuilder::from(scan_b)
        .alias("s").unwrap()
        .filter(c("s", "id").eq(lit("x1"))).unwrap()
        .project(vec![c("s", "id")]).unwrap()
        .build().unwrap();

    // Top = s JOIN J ON s.id = J.<b.id>  -- J's SECOND `id` (index 1)
    LogicalPlanBuilder::from(s)
        .join_detailed(
            j, JoinType::Inner,
            (vec![k("s", "id")], vec![k("b", "id")]),
            None, datafusion::common::NullEquality::NullEqualsNothing,
        ).unwrap().build().unwrap()
}

async fn run(dynamic: bool) -> (Arc<dyn datafusion::physical_plan::ExecutionPlan>, Vec<RecordBatch>) {
    let (_dir, ta, tb) = make_tables();
    let ctx = ctx(dynamic, &ta, &tb).await;
    let scan_a = ctx.table("ta").await.unwrap().logical_plan().clone();
    let scan_b = ctx.table("tb").await.unwrap().logical_plan().clone();
    let logical = build_logical(scan_a, scan_b);

    // NOTE: logical optimization is bypassed so the hand-built join tree is not
    // reordered; the physical optimizer (incl. dynamic-filter pushdown) still runs.
    let state = ctx.state();
    let mut phys = state.query_planner()
        .create_physical_plan(&logical, &state).await.unwrap();
    for rule in state.physical_optimizers() {
        phys = rule.optimize(phys, state.config_options()).unwrap();
    }
    let batches = datafusion::physical_plan::collect(
        Arc::clone(&phys), ctx.task_ctx(),
    ).await.unwrap();
    (phys, batches)
}

#[tokio::test]
async fn dynamic_filter_must_not_change_the_answer() {
    let (_, on) = run(true).await;
    let (_, off) = run(false).await;
    let rows = |b: &[RecordBatch]| b.iter().map(|x| x.num_rows()).sum::<usize>();
    assert_eq!(rows(&on), rows(&off),
        "a physical optimizer flag changed the answer: ON={} OFF={}",
        rows(&on), rows(&off));
}

if self.allowed_indices.contains(&col.index())
    && let Ok(new_index) = self.child_schema.index_of(col.name())
{
    Ok(Transformed::yes(Arc::new(Column::new(col.name(), new_index))))
}