better tests

Author: adriangb

datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs

@@ -160,7 +160,123 @@ fn test_pushdown_into_scan_with_config_options() {
 }
 
 #[tokio::test]
-async fn test_hashjoin_parent_filter_pushdown() {
+async fn test_dynamic_filter_pushdown_through_hash_join_with_topk() {
+    use datafusion_common::JoinType;
+    use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode};
+
+    // Create build side with limited values
+    let build_batches = vec![record_batch!(
+        ("a", Utf8, ["aa", "ab"]),
+        ("b", Utf8, ["ba", "bb"]),
+        ("c", Float64, [1.0, 2.0])
+    )
+    .unwrap()];
+    let build_side_schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Utf8, false),
+        Field::new("b", DataType::Utf8, false),
+        Field::new("c", DataType::Float64, false),
+    ]));
+    let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema))
+        .with_support(true)
+        .with_batches(build_batches)
+        .build();
+
+    // Create probe side with more values
+    let probe_batches = vec![record_batch!(
+        ("d", Utf8, ["aa", "ab", "ac", "ad"]),
+        ("e", Utf8, ["ba", "bb", "bc", "bd"]),
+        ("f", Float64, [1.0, 2.0, 3.0, 4.0])
+    )
+    .unwrap()];
+    let probe_side_schema = Arc::new(Schema::new(vec![
+        Field::new("d", DataType::Utf8, false),
+        Field::new("e", DataType::Utf8, false),
+        Field::new("f", DataType::Float64, false),
+    ]));
+    let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema))
+        .with_support(true)
+        .with_batches(probe_batches)
+        .build();
+
+    // Create HashJoinExec
+    let on = vec![(
+        col("a", &build_side_schema).unwrap(),
+        col("d", &probe_side_schema).unwrap(),
+    )];
+    let join = Arc::new(
+        HashJoinExec::try_new(
+            build_scan,
+            probe_scan,
+            on,
+            None,
+            &JoinType::Inner,
+            None,
+            PartitionMode::Partitioned,
+            datafusion_common::NullEquality::NullEqualsNothing,
+        )
+        .unwrap(),
+    );
+
+    let join_schema = join.schema();
+
+    // Finally let's add a SortExec on the outside to test pushdown of dynamic filters
+    let sort_expr =
+        PhysicalSortExpr::new(col("e", &join_schema).unwrap(), SortOptions::default());
+    let plan = Arc::new(
+        SortExec::new(LexOrdering::new(vec![sort_expr]).unwrap(), join)
+            .with_fetch(Some(2)),
+    ) as Arc<dyn ExecutionPlan>;
+
+    let mut config = ConfigOptions::default();
+    config.optimizer.enable_dynamic_filter_pushdown = true;
+    config.execution.parquet.pushdown_filters = true;
+
+    // Appy the FilterPushdown optimizer rule
+    let plan = FilterPushdown::new_post_optimization()
+        .optimize(Arc::clone(&plan), &config)
+        .unwrap();
+
+    // Test that filters are pushed down correctly to each side of the join
+    insta::assert_snapshot!(
+        format_plan_for_test(&plan),
+        @r"
+    - SortExec: TopK(fetch=2), expr=[e@4 ASC], preserve_partitioning=[false]
+    -   HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, d@0)]
+    -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=true
+    -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[d, e, f], file_type=test, pushdown_supported=true, predicate=DynamicFilterPhysicalExpr [ true ]
+    "
+    );
+
+    // Put some data through the plan to check that the filter is updated to reflect the TopK state
+    let session_ctx = SessionContext::new_with_config(SessionConfig::new());
+    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();
+    // Iterate one batch
+    stream.next().await.unwrap().unwrap();
+
+    // Test that filters are pushed down correctly to each side of the join
+    insta::assert_snapshot!(
+        format_plan_for_test(&plan),
+        @r"
+    - SortExec: TopK(fetch=2), expr=[e@4 ASC], preserve_partitioning=[false], filter=[e@4 IS NULL OR e@4 < bb]
+    -   HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, d@0)]
+    -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=true
+    -     DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[d, e, f], file_type=test, pushdown_supported=true, predicate=DynamicFilterPhysicalExpr [ e@1 IS NULL OR e@1 < bb ]
+    "
+    );
+}
+
+// Test both static and dynamic filter pushdown in HashJoinExec.
+// Note that static filter pushdown is rare: it should have already happened in the logical optimizer phase.
+// However users may manually construct plans that could result in a FilterExec -> HashJoinExec -> Scan setup.
+// Dynamic filters arise in cases such as nested inner joins or TopK -> HashJoinExec -> Scan setups.
+#[tokio::test]
+async fn test_static_filter_pushdown_through_hash_join() {
     use datafusion_common::JoinType;
     use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode};
 
@@ -283,7 +399,8 @@ async fn test_hashjoin_parent_filter_pushdown() {
 
     let join_schema = join.schema();
     let filter = col_lit_predicate("a", "aa", &join_schema);
-    let plan = Arc::new(FilterExec::try_new(filter, join).unwrap());
+    let plan =
+        Arc::new(FilterExec::try_new(filter, join).unwrap()) as Arc<dyn ExecutionPlan>;
 
     // Test that filters are NOT pushed down for left join
     insta::assert_snapshot!(

datafusion/physical-plan/src/joins/hash_join.rs

@@ -952,28 +952,22 @@ impl ExecutionPlan for HashJoinExec {
 
     fn gather_filters_for_pushdown(
         &self,
-        phase: FilterPushdownPhase,
+        _phase: FilterPushdownPhase,
         parent_filters: Vec<Arc<dyn PhysicalExpr>>,
         _config: &ConfigOptions,
     ) -> Result<FilterDescription> {
+        // Other types of joins can support *some* filters, but restrictions are complex and error prone.
+        // For now we don't support them.
+        // See the logical optimizer rules for more details: datafusion/optimizer/src/push_down_filter.rs
+        // See https://github.com/apache/datafusion/issues/16973 for tracking.
         if self.join_type != JoinType::Inner {
-            // Other types of joins can support *some* filters, but restrictions are complex and error prone.
-            // For now we don't support them.
-            // See the logical optimizer rules for more details: datafusion/optimizer/src/push_down_filter.rs
             return Ok(FilterDescription::all_unsupported(
                 &parent_filters,
                 &self.children(),
             ));
         }
-        if phase == FilterPushdownPhase::Pre {
-            FilterDescription::from_children(parent_filters, &self.children())
-        } else {
-            // TODO: push down our self filters to children
-            Ok(FilterDescription::all_unsupported(
-                &parent_filters,
-                &self.children(),
-            ))
-        }
+        FilterDescription::from_children(parent_filters, &self.children())
+        // TODO: push down our self filters to children in the post optimization phase
     }
 
     fn handle_child_pushdown_result(