wip

Author: jayshrivastava

datafusion/core/tests/physical_optimizer/filter_pushdown.rs

@@ -3422,6 +3422,176 @@ fn test_pushdown_through_aggregate_grouping_sets_with_reordered_input() {
     );
 }
 
+#[tokio::test]
+async fn test_hashjoin_dynamic_filter_pushdown_through_aggregate_with_reordered_input() {
+    // Reproduces the real-world scenario where a HashJoin's dynamic filter must
+    // push through an AggregateExec whose input columns are reordered by a
+    // ProjectionExec. The plan looks like:
+    //
+    //   HashJoinExec: on=[(bhandle@0, bhandle@0)]  ← creates DynFilter
+    //     DataSourceExec (build side)
+    //     AggregateExec: gby=[bhandle@1], aggr=[min(value)]  ← must pass filter
+    //       ProjectionExec: [value@1, bhandle@0]              ← reorders columns
+    //         DataSourceExec (probe side)                     ← should receive filter
+    //
+    // Without the fix, the aggregate blocks the dynamic filter because
+    // grouping expr Column("bhandle", 1) != filter Column("bhandle", 0).
+
+    use datafusion_common::JoinType;
+    use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode};
+
+    // Build side: simple lookup table with bhandle values
+    let build_batches = vec![record_batch!(
+        ("bhandle", Utf8, ["h1", "h2"])
+    )
+    .unwrap()];
+    let build_schema = Arc::new(Schema::new(vec![
+        Field::new("bhandle", DataType::Utf8, false),
+    ]));
+    let build_scan = TestScanBuilder::new(Arc::clone(&build_schema))
+        .with_support(true)
+        .with_batches(build_batches)
+        .build();
+
+    // Probe side: scan with (bhandle, value)
+    let probe_batches = vec![record_batch!(
+        ("bhandle", Utf8, ["h1", "h2", "h3", "h4"]),
+        ("value", Float64, [1.0, 2.0, 3.0, 4.0])
+    )
+    .unwrap()];
+    let probe_schema = Arc::new(Schema::new(vec![
+        Field::new("bhandle", DataType::Utf8, false),
+        Field::new("value", DataType::Float64, false),
+    ]));
+    let probe_scan = TestScanBuilder::new(Arc::clone(&probe_schema))
+        .with_support(true)
+        .with_batches(probe_batches)
+        .build();
+
+    // ProjectionExec reorders (bhandle, value) → (value, bhandle)
+    let reordered_schema = Arc::new(Schema::new(vec![
+        Field::new("value", DataType::Float64, false),
+        Field::new("bhandle", DataType::Utf8, false),
+    ]));
+    let projection = Arc::new(
+        ProjectionExec::try_new(
+            vec![
+                (col("value", &probe_schema).unwrap(), "value".to_string()),
+                (
+                    col("bhandle", &probe_schema).unwrap(),
+                    "bhandle".to_string(),
+                ),
+            ],
+            probe_scan,
+        )
+        .unwrap(),
+    );
+
+    // AggregateExec: GROUP BY bhandle@1, min(value@0)
+    let aggregate_expr = vec![
+        AggregateExprBuilder::new(
+            min_udaf(),
+            vec![col("value", &reordered_schema).unwrap()],
+        )
+        .schema(reordered_schema.clone())
+        .alias("min_value")
+        .build()
+        .map(Arc::new)
+        .unwrap(),
+    ];
+    let group_by = PhysicalGroupBy::new_single(vec![(
+        col("bhandle", &reordered_schema).unwrap(), // bhandle@1 in input
+        "bhandle".to_string(),
+    )]);
+
+    let aggregate = Arc::new(
+        AggregateExec::try_new(
+            AggregateMode::Single,
+            group_by,
+            aggregate_expr,
+            vec![None],
+            projection,
+            reordered_schema,
+        )
+        .unwrap(),
+    );
+    // Aggregate output schema: (bhandle@0, min_value@1)
+
+    // HashJoinExec: CollectLeft, join on bhandle
+    let agg_output_schema = aggregate.schema();
+    let plan = Arc::new(
+        HashJoinExec::try_new(
+            build_scan,
+            aggregate,
+            vec![(
+                col("bhandle", &build_schema).unwrap(),
+                col("bhandle", &agg_output_schema).unwrap(),
+            )],
+            None,
+            &JoinType::Inner,
+            None,
+            PartitionMode::CollectLeft,
+            datafusion_common::NullEquality::NullEqualsNothing,
+            false,
+        )
+        .unwrap(),
+    ) as Arc<dyn ExecutionPlan>;
+
+    // Post-optimization: the HashJoin's dynamic filter on bhandle should push
+    // through the aggregate and reach the probe-side DataSource.
+    insta::assert_snapshot!(
+        OptimizationTest::new(Arc::clone(&plan), FilterPushdown::new_post_optimization(), true),
+        @r"
+    OptimizationTest:
+      input:
+        - HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(bhandle@0, bhandle@0)]
+        -   DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[bhandle], file_type=test, pushdown_supported=true
+        -   AggregateExec: mode=Single, gby=[bhandle@1 as bhandle], aggr=[min_value]
+        -     ProjectionExec: expr=[value@1 as value, bhandle@0 as bhandle]
+        -       DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[bhandle, value], file_type=test, pushdown_supported=true
+      output:
+        Ok:
+          - HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(bhandle@0, bhandle@0)]
+          -   DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[bhandle], file_type=test, pushdown_supported=true
+          -   AggregateExec: mode=Single, gby=[bhandle@1 as bhandle], aggr=[min_value]
+          -     ProjectionExec: expr=[value@1 as value, bhandle@0 as bhandle]
+          -       DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[bhandle, value], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ empty ]
+    "
+    );
+
+    // Actually execute the plan to verify the dynamic filter is populated
+    let mut config = ConfigOptions::default();
+    config.execution.parquet.pushdown_filters = true;
+    config.optimizer.enable_dynamic_filter_pushdown = true;
+    config.optimizer.enable_aggregate_dynamic_filter_pushdown = true;
+    let plan = FilterPushdown::new_post_optimization()
+        .optimize(plan, &config)
+        .unwrap();
+
+    let session_config = SessionConfig::new().with_batch_size(10);
+    let session_ctx = SessionContext::new_with_config(session_config);
+    session_ctx.register_object_store(
+        ObjectStoreUrl::parse("test://").unwrap().as_ref(),
+        Arc::new(InMemory::new()),
+    );
+    let state = session_ctx.state();
+    let task_ctx = state.task_ctx();
+    let mut stream = plan.execute(0, Arc::clone(&task_ctx)).unwrap();
+    stream.next().await.unwrap().unwrap();
+
+    // After execution, the dynamic filter should be populated with bhandle values
+    insta::assert_snapshot!(
+        format!("{}", format_plan_for_test(&plan)),
+        @r"
+    - HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(bhandle@0, bhandle@0)]
+    -   DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[bhandle], file_type=test, pushdown_supported=true
+    -   AggregateExec: mode=Single, gby=[bhandle@1 as bhandle], aggr=[min_value]
+    -     ProjectionExec: expr=[value@1 as value, bhandle@0 as bhandle]
+    -       DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[bhandle, value], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ bhandle@0 >= h1 AND bhandle@0 <= h2 AND bhandle@0 IN (SET) ([h1, h2]) ]
+    "
+    );
+}
+
 #[test]
 fn test_pushdown_with_computed_grouping_key() {
     // Test filter pushdown with computed grouping expression