perf: Use agg DistinctHandling in join optimization

datafusion/core/tests/sql/joins.rs

@@ -27,6 +27,71 @@ use datafusion_sql::unparser::plan_to_sql;
 
 use super::*;
 
+#[tokio::test]
+async fn volatile_join_filter_preserves_evaluations_below_min() -> Result<()> {
+    use arrow::array::record_batch;
+    use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf};
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    // Request sort-merge joins; these small inputs remain in a single partition.
+    let mut config = SessionConfig::new().with_target_partitions(2);
+    config.options_mut().optimizer.prefer_hash_join = false;
+    let ctx = SessionContext::new_with_config(config);
+
+    // Return true, false, true, ... across successive input rows, including
+    // across batches. Both arguments keep the predicate at the join.
+    let evaluations = AtomicUsize::new(0);
+    ctx.register_udf(create_udf(
+        "alternating",
+        vec![DataType::Int32, DataType::Int32],
+        DataType::Boolean,
+        Volatility::Volatile,
+        Arc::new(move |args| {
+            let args = ColumnarValue::values_to_arrays(args)?;
+            let len = args[0].len();
+            let first = evaluations.fetch_add(len, Ordering::Relaxed);
+            let values = BooleanArray::from_iter(
+                (first..first + len).map(|i| Some(i.is_multiple_of(2))),
+            );
+            Ok(ColumnarValue::Array(Arc::new(values)))
+        }),
+    ));
+
+    ctx.register_batch(
+        "l",
+        record_batch!(("id", Int32, vec![1, 2]), ("x", Int32, vec![20, 10]))?,
+    )?;
+    ctx.register_batch(
+        "r",
+        record_batch!(("id", Int32, vec![1, 1, 2]), ("y", Int32, vec![0, 0, 0]))?,
+    )?;
+
+    // The inner join evaluates both pairs for id=1, so id=2 receives the
+    // third (true) result. A semi join would stop after the first match for
+    // id=1, give id=2 the second (false) result, and incorrectly return 20.
+    let df = ctx
+        .sql(
+            "SELECT MIN(l.x) AS minimum FROM l JOIN r \
+             ON l.id = r.id AND alternating(l.x, r.y)",
+        )
+        .await?;
+    let plan = df.create_physical_plan().await?;
+    let formatted = displayable(plan.as_ref()).indent(true).to_string();
+    assert_contains!(formatted, "SortMergeJoinExec: join_type=Inner");
+    let batches = collect(plan, ctx.task_ctx()).await?;
+    assert_batches_eq!(
+        [
+            "+---------+",
+            "| minimum |",
+            "+---------+",
+            "| 10      |",
+            "+---------+",
+        ],
+        &batches
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn join_change_in_planner() -> Result<()> {
     let config = SessionConfig::new().with_target_partitions(8);

datafusion/optimizer/src/eliminate_aggregate_distinct.rs

@@ -165,18 +165,13 @@ mod tests {
     use super::*;
     use crate::OptimizerContext;
     use crate::assert_optimized_plan_eq_snapshot;
+    use crate::test::udfs::DistinctHandlingTestUDAF;
     use crate::test::*;
 
     use crate::single_distinct_to_groupby::SingleDistinctToGroupBy;
-    use arrow::datatypes::DataType;
-    use datafusion_expr::function::AccumulatorArgs;
-    use datafusion_expr::{
-        Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder,
-        Signature, Volatility, col, lit,
-    };
+    use datafusion_expr::{AggregateUDF, ExprFunctionExt, LogicalPlanBuilder, col, lit};
     use datafusion_functions_aggregate::expr_fn::{bit_xor, max, min, sum};
 
-    use std::hash::{Hash, Hasher};
     use std::sync::Arc;
 
     macro_rules! assert_optimized_plan_equal {
@@ -196,66 +191,12 @@ mod tests {
         }};
     }
 
-    /// A user defined aggregate that reports the [`DistinctHandling`] it was
-    /// built with.
+    /// A user defined aggregate that reports the given [`DistinctHandling`].
     ///
     /// The tests above read the tag off built-in functions, which only covers
-    /// the variants those functions happen to carry. This one exercises the
-    /// public API a third-party function uses: an
-    /// [`AggregateUDFImpl::distinct_handling`] override.
-    #[derive(Debug, Clone, PartialEq, Eq)]
-    struct TaggedUdaf {
-        name: &'static str,
-        handling: DistinctHandling,
-        signature: Signature,
-    }
-
-    impl TaggedUdaf {
-        fn new(name: &'static str, handling: DistinctHandling) -> Self {
-            Self {
-                name,
-                handling,
-                signature: Signature::any(1, Volatility::Immutable),
-            }
-        }
-    }
-
-    /// Hashed by name, which identifies the function here. `DistinctHandling`
-    /// is not `Hash`, and `AggregateUDFImpl` requires one through `DynHash`.
-    impl Hash for TaggedUdaf {
-        fn hash<H: Hasher>(&self, state: &mut H) {
-            self.name.hash(state);
-            self.signature.hash(state);
-        }
-    }
-
-    impl AggregateUDFImpl for TaggedUdaf {
-        fn name(&self) -> &str {
-            self.name
-        }
-
-        fn signature(&self) -> &Signature {
-            &self.signature
-        }
-
-        fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-            Ok(DataType::UInt32)
-        }
-
-        fn accumulator(
-            &self,
-            _acc_args: AccumulatorArgs,
-        ) -> Result<Box<dyn Accumulator>> {
-            unimplemented!("the rule only rewrites the logical plan")
-        }
-
-        fn distinct_handling(&self) -> DistinctHandling {
-            self.handling
-        }
-    }
-
+    /// the variants those functions happen to carry.
     fn tagged(name: &'static str, handling: DistinctHandling) -> AggregateUDF {
-        AggregateUDF::from(TaggedUdaf::new(name, handling))
+        AggregateUDF::from(DistinctHandlingTestUDAF::new(name, handling))
     }
 
     /// `min(DISTINCT b)` loses the flag but keeps its column name.