diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 1c23134e778f5..da4f2d2565655 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -472,9 +472,14 @@ pub fn aggregate_functional_dependencies( // The loop below only re-expresses input dependencies. Skip it when the // input has none. The GROUP BY-key dependency below always runs. if !func_dependencies.is_empty() { - // Compute once: this does not change in the loop. - let existing_target_indices = - get_target_functional_dependencies(aggr_input_schema, group_by_input_indices); + // Compare the columns covered by every dependency, whatever its NULL + // semantics: leaving out the ones that do not hold across NULLs would + // hide GROUP BY expressions this determinant does not cover and + // upgrade its mode to `Single` wrongly. + let existing_target_indices = get_target_indices_from_dependencies( + func_dependencies.iter(), + group_by_input_indices, + ); for FunctionalDependence { source_indices, nullable, @@ -500,8 +505,8 @@ pub fn aggregate_functional_dependencies( new_source_input_indices.push(Some(*input_idx)); } } - let new_target_indices = get_target_functional_dependencies( - aggr_input_schema, + let new_target_indices = get_target_indices_from_dependencies( + func_dependencies.iter(), &new_source_input_indices, ); let mode = if existing_target_indices == new_target_indices @@ -572,6 +577,12 @@ pub fn aggregate_functional_dependencies( /// Returns target indices, for the determinant keys that are inside /// group by expressions. /// +/// Only dependencies that hold when NULL determinant values are grouped +/// together are considered (see +/// [`FunctionalDependence::is_valid_across_nulls`]): a nullable `UNIQUE` +/// column permits multiple NULL rows that may differ in the other columns, +/// so it does not determine them across the single NULL group. +/// /// `group_by_input_indices` holds, for each GROUP BY expression, the index of /// the `schema` field it references, or `None` for a computed expression. pub fn get_target_functional_dependencies( @@ -582,21 +593,27 @@ pub fn get_target_functional_dependencies( if dependencies.is_empty() { return None; } + let dependencies = dependencies + .iter() + .filter(|dependency| dependency.is_valid_across_nulls(schema)); + get_target_indices_from_dependencies(dependencies, group_by_input_indices) +} + +fn get_target_indices_from_dependencies<'a>( + dependencies: impl IntoIterator, + group_by_input_indices: &[Option], +) -> Option> { let mut combined_target_indices = HashSet::new(); - for FunctionalDependence { - source_indices, - target_indices, - .. - } in &dependencies.deps - { + for dependency in dependencies { // If the GROUP BY expression contains a determinant key, we can use // the associated fields after aggregation even if they are not part // of the GROUP BY expression. - if source_indices + if dependency + .source_indices .iter() .all(|source_idx| group_by_input_indices.contains(&Some(*source_idx))) { - combined_target_indices.extend(target_indices.iter()); + combined_target_indices.extend(dependency.target_indices.iter()); } } (!combined_target_indices.is_empty()).then_some({ @@ -623,22 +640,21 @@ pub fn get_required_group_by_exprs_indices( .collect::>>()?; groupby_expr_indices.sort_unstable(); - for FunctionalDependence { - source_indices, - target_indices, - .. - } in &dependencies.deps - { - if source_indices + for dependency in &dependencies.deps { + // If all source indices are among GROUP BY expression indices, we + // can remove target indices from GROUP BY expression indices and + // use source indices instead. The dependency must also hold across + // NULL keys, which GROUP BY collapses into a single group. + if dependency + .source_indices .iter() .all(|source_idx| groupby_expr_indices.contains(source_idx)) + && dependency.is_valid_across_nulls(schema) { - // If all source indices are among GROUP BY expression indices, we - // can remove target indices from GROUP BY expression indices and - // use source indices instead. - groupby_expr_indices = set_difference(&groupby_expr_indices, target_indices); groupby_expr_indices = - merge_and_order_indices(groupby_expr_indices, source_indices); + set_difference(&groupby_expr_indices, &dependency.target_indices); + groupby_expr_indices = + merge_and_order_indices(groupby_expr_indices, &dependency.source_indices); } } groupby_expr_indices diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index aa81df22554c2..f443043d38835 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -3138,6 +3138,63 @@ mod tests { Ok(()) } + #[test] + fn plan_builder_aggregate_nullable_unique_adds_no_implicit_group_by_exprs() + -> Result<()> { + // A nullable UNIQUE column permits multiple NULL rows that may differ + // in the other columns, so it does not determine them across the + // single NULL group and they must not be added to the GROUP BY. + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("state", DataType::Utf8, true), + Field::new("salary", DataType::Int32, true), + ]); + let constraints = Constraints::new_unverified(vec![Constraint::Unique(vec![0])]); + let table_source = table_source_with_constraints(&schema, constraints); + + let options = + LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true); + let plan = LogicalPlanBuilder::scan("employee_csv", table_source, None)? + .with_options(options) + .aggregate(vec![col("id")], vec![sum(col("salary"))])? + .build()?; + + assert_snapshot!(plan, @r" + Aggregate: groupBy=[[employee_csv.id]], aggr=[[sum(employee_csv.salary)]] + TableScan: employee_csv + "); + + Ok(()) + } + + #[test] + fn plan_builder_aggregate_not_null_unique_implicit_group_by_exprs() -> Result<()> { + // A UNIQUE column declared NOT NULL has no NULL rows, so it determines + // the other columns like a PRIMARY KEY and they are added to the + // GROUP BY. + let schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("state", DataType::Utf8, true), + Field::new("salary", DataType::Int32, true), + ]); + let constraints = Constraints::new_unverified(vec![Constraint::Unique(vec![0])]); + let table_source = table_source_with_constraints(&schema, constraints); + + let options = + LogicalPlanBuilderOptions::new().with_add_implicit_group_by_exprs(true); + let plan = LogicalPlanBuilder::scan("employee_csv", table_source, None)? + .with_options(options) + .aggregate(vec![col("id")], vec![sum(col("salary"))])? + .build()?; + + assert_snapshot!(plan, @r" + Aggregate: groupBy=[[employee_csv.id, employee_csv.state, employee_csv.salary]], aggr=[[sum(employee_csv.salary)]] + TableScan: employee_csv + "); + + Ok(()) + } + #[test] fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> { // https://github.com/apache/datafusion/issues/23812 diff --git a/datafusion/optimizer/src/replace_distinct_aggregate.rs b/datafusion/optimizer/src/replace_distinct_aggregate.rs index 5400a9feeccdf..42114f84ca132 100644 --- a/datafusion/optimizer/src/replace_distinct_aggregate.rs +++ b/datafusion/optimizer/src/replace_distinct_aggregate.rs @@ -102,8 +102,11 @@ impl OptimizerRule for ReplaceDistinctWithAggregate { // is a no-op and we can simply remove it. The dependency mode // must be `Single`: a `Multi` dependence (e.g. a former key // downgraded by a join) means equal rows may occur multiple - // times, so the DISTINCT still has work to do. + // times, so the DISTINCT still has work to do. Uniqueness + // must also hold across NULLs: a nullable UNIQUE key can + // contain multiple NULL rows that DISTINCT must collapse. if dep.mode == Dependency::Single + && dep.is_valid_across_nulls(input.schema()) && dep.source_indices.len() >= field_count && dep.source_indices[..field_count] .iter() diff --git a/datafusion/sqllogictest/test_files/functional_dependencies.slt b/datafusion/sqllogictest/test_files/functional_dependencies.slt index 60ba1f00e6cca..99f1d001397c1 100644 --- a/datafusion/sqllogictest/test_files/functional_dependencies.slt +++ b/datafusion/sqllogictest/test_files/functional_dependencies.slt @@ -59,6 +59,18 @@ SELECT x, y FROM t_uniq; NULL 1 NULL 2 +# `x` is UNIQUE but declared NOT NULL, so there are no NULL rows to worry about. +statement ok +CREATE TABLE t_uniq_not_null (x INT NOT NULL UNIQUE, y INT) AS VALUES (1, 10), (2, 20); + +# `t_right_pk` is the NULL-padded side of `t_left LEFT JOIN t_right_pk`: +# `k = 3` and `k = 4` match nothing, so `p.a` and `p.b` are both NULL there. +statement ok +CREATE TABLE t_left (k INT) AS VALUES (1), (3), (4); + +statement ok +CREATE TABLE t_right_pk (a INT, b INT, PRIMARY KEY (a)) AS VALUES (1, 10), (2, 20); + # 1.1 PRIMARY KEY: rows are unique; the DISTINCT is removed and no # Aggregate appears in the plan. @@ -69,21 +81,18 @@ logical_plan TableScan: t_pk projection=[x] # 1.2 Nullable UNIQUE: the DISTINCT must be KEPT. UNIQUE allows several NULL # rows, but DISTINCT treats NULLs as equal and has to collapse them into one. -# -# BUG: the DISTINCT is removed and both NULL rows are returned. -# Expected: `1`, `NULL`. -# Issue: https://github.com/apache/datafusion/issues/23634 query I SELECT DISTINCT x FROM t_uniq ORDER BY x NULLS LAST; ---- 1 NULL -NULL query TT EXPLAIN SELECT DISTINCT x FROM t_uniq; ---- -logical_plan TableScan: t_uniq projection=[x] +logical_plan +01)Aggregate: groupBy=[[t_uniq.x]], aggr=[[]] +02)--TableScan: t_uniq projection=[x] # 1.3 A PRIMARY KEY downgraded to a non-unique dependency by a LEFT JOIN # so the DISTINCT must be KEPT. @@ -124,6 +133,61 @@ logical_plan 02)--TableScan: t_uniq projection=[x] +# 1.5 GROUP BY x, y keeps separate groups for the two NULL rows, so the +# projected x is not unique. DISTINCT must still collapse the NULL values. +query I +SELECT DISTINCT x FROM t_uniq GROUP BY x, y ORDER BY x NULLS LAST; +---- +1 +NULL + +# Aliases and subqueries must not turn the nullable dependency into a proof +# that the projected values are unique. +query I +SELECT DISTINCT x AS k FROM t_uniq GROUP BY x, y ORDER BY k NULLS LAST; +---- +1 +NULL + +query I +SELECT DISTINCT x FROM (SELECT x, y FROM t_uniq GROUP BY x, y) g ORDER BY x NULLS LAST; +---- +1 +NULL + +# An outer aggregate must count distinct x values, not the inner (x, y) groups. +query I +SELECT count(*) FROM (SELECT DISTINCT x FROM t_uniq GROUP BY x, y) g; +---- +2 + +# 1.6 A composite UNIQUE key is also unsafe when only one component is nullable. +statement ok +CREATE TABLE t_uniq_composite (x INT NOT NULL, y INT, z INT, UNIQUE (x, y)) +AS VALUES (1, NULL, 10), (1, NULL, 20), (1, 1, 30); + +query II +SELECT DISTINCT x, y FROM t_uniq_composite GROUP BY x, y, z ORDER BY x, y NULLS LAST; +---- +1 1 +1 NULL + +statement ok +drop table t_uniq_composite; + +# 1.7 NOT NULL UNIQUE still guarantees uniqueness, so DISTINCT is redundant. +query I +SELECT DISTINCT x FROM t_uniq_not_null ORDER BY x; +---- +1 +2 + +query TT +EXPLAIN SELECT DISTINCT x FROM t_uniq_not_null; +---- +logical_plan TableScan: t_uniq_not_null projection=[x] + + # 2.1 PRIMARY KEY: `x` determines `y`, so `ORDER BY x, y` is equivalent to # `ORDER BY x` and the `y` key is dropped from the plan. query TT @@ -151,9 +215,6 @@ logical_plan 02)--TableScan: t_uniq projection=[x, y] # 2.3 A UNIQUE determinant remains usable when its source column is NOT NULL. -statement ok -CREATE TABLE t_uniq_not_null (x INT NOT NULL UNIQUE, y INT) AS VALUES (1, 10), (2, 20); - query TT EXPLAIN SELECT x, y FROM t_uniq_not_null ORDER BY x, y; ---- @@ -161,9 +222,6 @@ logical_plan 01)Sort: t_uniq_not_null.x ASC NULLS LAST 02)--TableScan: t_uniq_not_null projection=[x, y] -statement ok -drop table t_uniq_not_null; - # 2.4 After `GROUP BY x` the `x` does determine `cnt`, so can drop `cnt` from sort query TT EXPLAIN SELECT x, cnt FROM (SELECT x, count(*) AS cnt FROM t_uniq GROUP BY x) ORDER BY x, cnt; @@ -177,12 +235,6 @@ logical_plan # 2.5 A PRIMARY KEY on the NULL-padded side of a LEFT JOIN: every new NULL in # `p.a` comes from a padded row, where `p.b` is NULL too. So `p.a` still # determines `p.b` and the `p.b` sort key can be dropped. -statement ok -CREATE TABLE t_left (k INT) AS VALUES (1), (3), (4); - -statement ok -CREATE TABLE t_right_pk (a INT, b INT, PRIMARY KEY (a)) AS VALUES (1, 10), (2, 20); - query II SELECT p.a, p.b FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a ORDER BY p.a, p.b DESC; ---- @@ -201,12 +253,6 @@ logical_plan 05)----SubqueryAlias: p 06)------TableScan: t_right_pk projection=[a, b] -statement ok -drop table t_left; - -statement ok -drop table t_right_pk; - # 2.6 After `GROUP BY x, y` the pair `(x, y)` is unique (NULLs included), so # it determines `cnt` even though the nullable UNIQUE `x` alone does not. query III @@ -236,27 +282,54 @@ logical_plan 02)--TableScan: t_pk projection=[x] # 3.2 Nullable UNIQUE: grouping by `x, y` is NOT the same as grouping by -# `x` -- two NULL rows differ in `y` and belong in separate groups. -# -# BUG: `y` is dropped from the GROUP BY and the two NULL groups are merged, -# so one row goes missing. -# Expected: `1`, `NULL`, `NULL` (three rows). -# Issue: https://github.com/apache/datafusion/issues/23819 +# `x` -- two NULL rows differ in `y` and belong in separate groups, so the +# `y` grouping key must be kept. query I rowsort SELECT x FROM t_uniq GROUP BY x, y; ---- 1 NULL +NULL query TT EXPLAIN SELECT x FROM t_uniq GROUP BY x, y; ---- logical_plan -01)Aggregate: groupBy=[[t_uniq.x]], aggr=[[]] -02)--TableScan: t_uniq projection=[x] +01)Projection: t_uniq.x +02)--Aggregate: groupBy=[[t_uniq.x, t_uniq.y]], aggr=[[]] +03)----TableScan: t_uniq projection=[x, y] -# 3.3 The same grouping, but with `y` selected so the parent needs it: no -# column can be dropped and the answer is right. +# 3.3 A UNIQUE determinant remains usable when its source column is NOT NULL. +query TT +EXPLAIN SELECT x FROM t_uniq_not_null GROUP BY x, y; +---- +logical_plan +01)Aggregate: groupBy=[[t_uniq_not_null.x]], aggr=[[]] +02)--TableScan: t_uniq_not_null projection=[x] + +# 3.4 A PRIMARY KEY on the NULL-padded side of a LEFT JOIN: every new NULL in +# `p.a` comes from a padded row, where `p.b` is NULL too. So `p.a` still +# determines `p.b` and the `p.b` grouping key can be dropped. +query I +SELECT p.a FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a GROUP BY p.a, p.b ORDER BY p.a; +---- +1 +NULL + +query TT +EXPLAIN SELECT p.a FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a GROUP BY p.a, p.b; +---- +logical_plan +01)Aggregate: groupBy=[[p.a]], aggr=[[]] +02)--Projection: p.a +03)----Left Join: l.k = p.a +04)------SubqueryAlias: l +05)--------TableScan: t_left projection=[k] +06)------SubqueryAlias: p +07)--------TableScan: t_right_pk projection=[a] + +# 3.5 The same grouping as 3.2, but with `y` selected so the parent needs it: +# no column can be dropped and the answer is right. query II rowsort SELECT x, y FROM t_uniq GROUP BY x, y; ---- @@ -280,26 +353,46 @@ logical_plan 02)--TableScan: t_pk projection=[x, y] # 4.2 Nullable UNIQUE: `x` does NOT determine `y`, so there is no -# well-defined `y` for the `x = NULL` group. -# -# BUG: `y` is appended to the GROUP BY anyway, so `GROUP BY x` returns TWO -# rows for `x = NULL`. -# Expected: one row per distinct `x` (or a planning error -- postgres -# rejects this query, and accepts the 4.1 PRIMARY KEY form). -# Issue: https://github.com/apache/datafusion/issues/23820 -query II rowsort +# well-defined `y` for the `x = NULL` group and `y` is not added to the +# GROUP BY. The query is rejected, as postgres does (it accepts the 4.1 +# PRIMARY KEY form). +query error column "t_uniq.y" must appear in the GROUP BY clause or must be part of an aggregate function SELECT x, y FROM t_uniq GROUP BY x; + +# 4.3 A UNIQUE determinant remains usable when its source column is NOT NULL. +query II rowsort +SELECT x, y FROM t_uniq_not_null GROUP BY x; ---- -1 3 -NULL 1 -NULL 2 +1 10 +2 20 query TT -EXPLAIN SELECT x, y FROM t_uniq GROUP BY x; +EXPLAIN SELECT x, y FROM t_uniq_not_null GROUP BY x; ---- logical_plan -01)Aggregate: groupBy=[[t_uniq.x, t_uniq.y]], aggr=[[]] -02)--TableScan: t_uniq projection=[x, y] +01)Aggregate: groupBy=[[t_uniq_not_null.x, t_uniq_not_null.y]], aggr=[[]] +02)--TableScan: t_uniq_not_null projection=[x, y] + +# 4.4 A PRIMARY KEY on the NULL-padded side of a LEFT JOIN still determines +# its other columns (padded rows are NULL in all of them), so `p.b` can be +# selected when grouping by `p.a` alone. +query II +SELECT p.a, p.b FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a GROUP BY p.a ORDER BY p.a; +---- +1 10 +NULL NULL + +query TT +EXPLAIN SELECT p.a, p.b FROM t_left l LEFT JOIN t_right_pk p ON l.k = p.a GROUP BY p.a; +---- +logical_plan +01)Aggregate: groupBy=[[p.a, p.b]], aggr=[[]] +02)--Projection: p.a, p.b +03)----Left Join: l.k = p.a +04)------SubqueryAlias: l +05)--------TableScan: t_left projection=[k] +06)------SubqueryAlias: p +07)--------TableScan: t_right_pk projection=[a, b] statement ok @@ -563,5 +656,14 @@ drop table t_pk; statement ok drop table t_uniq; +statement ok +drop table t_uniq_not_null; + +statement ok +drop table t_left; + +statement ok +drop table t_right_pk; + statement ok RESET datafusion.explain.logical_plan_only; diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 64c6b04d41b36..58f0ea392279c 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -3668,8 +3668,8 @@ ORDER BY r.sn query IRR SELECT l.sn, l.amount, SUM(l.amount) FROM (SELECT * - FROM sales_global_with_unique as l - LEFT SEMI JOIN sales_global_with_unique as r + FROM sales_global_with_pk as l + LEFT SEMI JOIN sales_global_with_pk as r ON l.amount >= r.amount + 10) GROUP BY l.sn ORDER BY l.sn @@ -3678,20 +3678,48 @@ ORDER BY l.sn 2 75 75 3 200 200 4 100 100 -NULL 100 100 # Similarly, left anti join should propagate constraint of left side as is. query IRR SELECT l.sn, l.amount, SUM(l.amount) FROM (SELECT * - FROM sales_global_with_unique as l - LEFT ANTI JOIN sales_global_with_unique as r + FROM sales_global_with_pk as l + LEFT ANTI JOIN sales_global_with_pk as r ON l.amount >= r.amount + 10) GROUP BY l.sn ORDER BY l.sn ---- 0 30 30 +# The propagated unique key of the left side can contain NULL, so it does not +# determine `l.amount` for the `l.sn = NULL` group and the query is rejected, +# just as it would be without the join. +statement error DataFusion error: Error during planning: Column in SELECT must be in GROUP BY or an aggregate function: While expanding wildcard, column "l\.amount" must appear in the GROUP BY clause or must be part of an aggregate function, currently only "l\.sn, sum\(l\.amount\)" appears in the SELECT clause satisfies this requirement +SELECT l.sn, l.amount, SUM(l.amount) + FROM (SELECT * + FROM sales_global_with_unique as l + LEFT SEMI JOIN sales_global_with_unique as r + ON l.amount >= r.amount + 10) + GROUP BY l.sn +ORDER BY l.sn + +# Grouping by `l.amount` explicitly is accepted, and the `l.sn = NULL` row +# comes back as its own group. +query IRR +SELECT l.sn, l.amount, SUM(l.amount) + FROM (SELECT * + FROM sales_global_with_unique as l + LEFT SEMI JOIN sales_global_with_unique as r + ON l.amount >= r.amount + 10) + GROUP BY l.sn, l.amount +ORDER BY l.sn +---- +1 50 50 +2 75 75 +3 200 200 +4 100 100 +NULL 100 100 + # Should support grouping by list column query ?I SELECT column1, COUNT(*) as column2 FROM (VALUES (['a', 'b'], 1), (['c', 'd', 'e'], 2), (['a', 'b'], 3)) as values0 GROUP BY column1 ORDER BY column2; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 2ec01c99479c4..7b6bc6a99e33c 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -986,6 +986,28 @@ let dep = FunctionalDependence::new(source_indices, target_indices, nullable) .with_null_equality(NullEquality::NullEqualsNull); ``` +### Nullable `UNIQUE` constraints no longer add or remove `GROUP BY` expressions + +`get_target_functional_dependencies` and `get_required_group_by_exprs_indices` +now ignore a functional dependency that does not hold across NULL determinant +values (see `FunctionalDependence::is_valid_across_nulls`), such as one derived +from a nullable `UNIQUE` constraint, because `GROUP BY` places all NULL keys in +a single group. + +As a result, a column that is only determined by a nullable `UNIQUE` key can no +longer be selected without appearing in the `GROUP BY`. This matches PostgreSQL, +which accepts such a query only for a `PRIMARY KEY`: + +```sql +CREATE TABLE t (x INT UNIQUE, y INT); + +-- Now fails: Column in SELECT must be in GROUP BY or an aggregate function +SELECT x, y FROM t GROUP BY x; + +-- Works: add `y` to the GROUP BY, aggregate it, or declare `x` NOT NULL +SELECT x, y FROM t GROUP BY x, y; +``` + ### Functional dependency helpers take field indices instead of names `aggregate_functional_dependencies`, `get_target_functional_dependencies`,