neilconway commented on code in PR #25385:
URL: https://github.com/apache/datafusion/pull/25385#discussion_r4053733463
##########
datafusion/optimizer/src/eliminate_join.rs:
##########
@@ -56,11 +57,12 @@
//! set across its two inputs.
//! * `duplicate_insensitive` — whether emitting each row once instead of many
//! times will not change the output. A duplicate-collapsing node (e.g.,
-//! DISTINCT, GROUP BY with no aggregate functions, or the existence side of
a
-//! semi/anti/mark join) sets it `true` for its subtree, and it propagates
-//! downward until a node that makes the row count observable again (a
`LIMIT`,
-//! a top-N sort, ...) clears it. It is therefore fixed by the nearest such
-//! node, not by the whole ancestor chain: a collapsing node shields its
subtree,
+//! DISTINCT, an `Aggregate` plan node whose aggregate expressions all ignore
+//! duplicate input rows, or the existence side of a semi/anti/mark join)
sets
+//! it `true` for its subtree, and it propagates downward until a node that
+//! makes the row count observable again (a `LIMIT`, a top-N sort, a volatile
+//! expression, ...) clears it. It is therefore fixed by the nearest such
node,
+//! not by the whole ancestor chain: a collapsing node shields its subtree,
Review Comment:
Thanks! Fixed.
##########
datafusion/optimizer/src/eliminate_join.rs:
##########
@@ -769,6 +806,438 @@ mod tests {
")
}
+ #[test]
+ fn insensitive_aggregates_enable_semi_joins() -> Result<()> {
+ for column in ["l.x", "r.x"] {
+ let aggr_expr = vec![
+ min(col(column)).alias("minimum"),
+ max(col(column)).distinct().build()?,
+ ];
+ // Both global and grouped aggregates ignore duplicate input rows.
+ for group_expr in [vec![], vec![col(column)]] {
+ let plan = left_join_right()?
+ .aggregate(group_expr, aggr_expr.clone())?
+ .build()?;
+ let result =
+ EliminateJoin::new().rewrite(plan,
&OptimizerContext::new())?;
+ assert!(result.transformed);
+ let LogicalPlan::Aggregate(aggregate) = result.data else {
+ panic!("expected aggregate");
+ };
+ assert_eq!(aggregate.aggr_expr, aggr_expr);
+ let LogicalPlan::Join(join) = aggregate.input.as_ref() else {
+ panic!("expected join");
+ };
+ assert_eq!(
+ join.join_type,
+ if column == "l.x" {
+ JoinType::LeftSemi
+ } else {
+ JoinType::RightSemi
+ }
+ );
+ }
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn global_min_removes_unused_outer_join() -> Result<()> {
+ for (join_type, column, table) in
+ [(JoinType::Left, "l.x", "l"), (JoinType::Right, "r.x", "r")]
+ {
+ let left = scan("l", &test_schema(), Constraints::default())?;
+ let right = scan("r", &test_schema(), Constraints::default())?;
+ let plan = LogicalPlanBuilder::from(left)
+ .join(right, join_type, (vec!["l.id"], vec!["r.id"]), None)?
+ .aggregate(Vec::<Expr>::new(), vec![min(col(column))])?
+ .build()?;
+ let optimized = EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .data;
+ let expected = LogicalPlanBuilder::from(scan(
+ table,
+ &test_schema(),
+ Constraints::default(),
+ )?)
+ .aggregate(Vec::<Expr>::new(), vec![min(col(column))])?
+ .build()?;
+ assert_eq!(optimized, expected);
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn distinct_sensitive_aggregates_enable_semi_joins() -> Result<()> {
+ // A `Sensitive` function called with DISTINCT deduplicates its own
+ // input, so it cannot observe rows repeated by the join.
+ for aggr_expr in [
+ vec![count_distinct(col("l.x"))],
+ vec![count_distinct(col("l.x")), count_distinct(col("l.y"))],
+ vec![
+ min(col("l.x")),
+ count(col("l.x"))
+ .distinct()
+ .filter(col("l.y").gt(lit(0)))
+ .build()?,
+ ],
+ ] {
+ let plan = left_join_right()?
+ .aggregate(vec![col("l.id")], aggr_expr)?
+ .build()?;
+ let result = EliminateJoin::new().rewrite(plan,
&OptimizerContext::new())?;
+ assert!(result.transformed);
+ let LogicalPlan::Aggregate(aggregate) = result.data else {
+ panic!("expected aggregate");
+ };
+ let LogicalPlan::Join(join) = aggregate.input.as_ref() else {
+ panic!("expected join");
+ };
+ assert_eq!(join.join_type, JoinType::LeftSemi);
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn duplicate_sensitive_aggregates_block_rewrite() -> Result<()> {
+ // One aggregate that observes repeated rows keeps the join, even
+ // beside aggregates that do not. DISTINCT does not qualify an
+ // `Unsupported` function: its accumulator does not deduplicate, and
+ // may silently compute the non-distinct answer.
+ for sensitive in [
+ count(col("l.x")),
+ stddev(col("l.x")).distinct().build()?,
+ corr(col("l.x"), col("l.y")).distinct().build()?,
+ regr_count(col("l.x"), col("l.y")).distinct().build()?,
+ ] {
+ let plan = left_join_right()?
+ .aggregate(
+ Vec::<Expr>::new(),
+ vec![min(col("l.x")), count_distinct(col("l.x")),
sensitive],
+ )?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn sensitive_aggregate_blocks_insensitive_ancestor() -> Result<()> {
+ let plan = left_join_right()?
+ .aggregate(vec![col("l.x")], vec![count(col("l.id")).alias("n")])?
+ .aggregate(Vec::<Expr>::new(), vec![min(col("n"))])?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn subquery_aggregate_argument_blocks_rewrite() -> Result<()> {
+ // Expr's usual volatility check does not descend into a subquery plan.
+ let volatile = ScalarUDF::from(
+ PlacementTestUDF::new().with_volatility(Volatility::Volatile),
+ )
+ .call(vec![lit(1)]);
+ let subquery = LogicalPlanBuilder::empty(true)
+ .project(vec![volatile])?
+ .build()?;
+ let plan = left_join_right()?
+ .aggregate(
+ vec![col("l.x")],
+ vec![min(scalar_subquery(Arc::new(subquery)))],
+ )?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ Ok(())
+ }
+
+ #[test]
+ fn aggregate_filter_and_ordering_keep_columns_live() -> Result<()> {
+ for aggr in [
+ min(col("l.x")).filter(col("r.y").gt(lit(0))).build()?,
+ min(col("l.x"))
+ .order_by(vec![col("r.y").sort(true, false)])
+ .build()?,
+ ] {
+ let plan = left_join_right()?
+ .aggregate(Vec::<Expr>::new(), vec![aggr])?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ }
+
+ let plan = left_join_right()?
+ .aggregate(
+ Vec::<Expr>::new(),
+ vec![min(col("l.x")).filter(col("l.y").gt(lit(0))).build()?],
+ )?
+ .build()?;
+ assert_optimized_plan_equal!(plan, @r"
+ Aggregate: groupBy=[[]], aggr=[[min(l.x) FILTER (WHERE l.y >
Int32(0))]]
+ LeftSemi Join: l.id = r.id
+ TableScan: l
+ TableScan: r
+ ")
+ }
+
+ fn volatile_expr() -> Expr {
+
ScalarUDF::from(PlacementTestUDF::new().with_volatility(Volatility::Volatile))
+ .call(vec![col("l.x")])
+ }
+
+ #[test]
+ fn volatile_aggregate_expressions_block_rewrite() -> Result<()> {
+ for (group_expr, aggr) in [
+ (vec![], min(volatile_expr())),
+ (vec![volatile_expr()], min(col("l.x"))),
+ (
+ vec![],
+ min(col("l.x"))
+ .filter(volatile_expr().gt(lit(0_u32)))
+ .build()?,
+ ),
+ (
+ vec![],
+ min(col("l.x"))
+ .order_by(vec![volatile_expr().sort(true, false)])
+ .build()?,
+ ),
+ ] {
+ let plan = left_join_right()?
+ .aggregate(group_expr, vec![aggr])?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ }
+ Ok(())
+ }
+
+ #[test]
+ fn volatile_intervening_expressions_block_rewrite() -> Result<()> {
+ for input in [
+ left_join_right()?.project(vec![col("l.x"),
volatile_expr().alias("v")])?,
+ left_join_right()?.filter(volatile_expr().gt(lit(0_u32)))?,
+ left_join_right()?.sort(vec![volatile_expr().sort(true, false)])?,
+ ] {
+ let plan = input
+ .aggregate(Vec::<Expr>::new(), vec![min(col("l.x"))])?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ }
+ Ok(())
+ }
Review Comment:
Good catch, fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
