adriangb commented on code in PR #26150:
URL: https://github.com/apache/datafusion/pull/26150#discussion_r4236255834
##########
datafusion/sqllogictest/test_files/functional_dependencies.slt:
##########
@@ -358,6 +358,89 @@ drop table t_null;
statement ok
drop table t_probe;
+
+# 6. Only a column reference carries a dependency. `CAST(x AS INT)` is named
+# `t_cast.x`, like the PRIMARY KEY column, but it does not determine `y`:
+# 1.1 and 1.2 both cast to 1.
+statement ok
+CREATE TABLE t_cast (x DOUBLE, y VARCHAR, PRIMARY KEY (x)) AS VALUES (1.1,
'b'), (1.2, 'a'), (2.1, 'c');
+
+# 6.1 ORDER BY: `y` is the tie-breaker within `CAST(x AS INT) = 1` and is kept.
+query RT
+SELECT x, y FROM t_cast ORDER BY CAST(x AS INT), y;
+----
+1.2 a
+1.1 b
+2.1 c
+
+query TT
+EXPLAIN SELECT x, y FROM t_cast ORDER BY CAST(x AS INT), y;
+----
+logical_plan
+01)Sort: CAST(t_cast.x AS Int32) ASC NULLS LAST, t_cast.y ASC NULLS LAST
+02)--TableScan: t_cast projection=[x, y]
+
+# 6.2 GROUP BY: `y` is not selected but still splits the groups.
+query II rowsort
+SELECT CAST(x AS INT) k, count(*) n FROM t_cast GROUP BY CAST(x AS INT), y;
+----
+1 1
+1 1
+2 1
+
+# 6.3 GROUP BY output: the dependency `k -> y` does not exist, so `y` still
+# orders the rows within `k = 1`.
+query IT
+SELECT CAST(x AS INT) k, y FROM t_cast GROUP BY CAST(x AS INT), y ORDER BY k,
y;
+----
+1 a
+1 b
+2 c
+
+# 6.4 `y` is not determined by the GROUP BY expression, so it can't be
selected.
+query error DataFusion error: Error during planning: Column in SELECT must be
in GROUP BY or an aggregate function
+SELECT y, count(*) FROM t_cast GROUP BY CAST(x AS INT);
+
+statement ok
+drop table t_cast;
+
+# 6.5 Only the key column itself carries the dependency, even when the cast
+# keeps distinct values distinct, like INT to BIGINT. PostgreSQL and DuckDB
+# reject this too.
+statement ok
+CREATE TABLE t_wide (x INT, y VARCHAR, PRIMARY KEY (x)) AS VALUES (1, 'b'),
(2, 'a');
+
+query error DataFusion error: Error during planning: Column in SELECT must be
in GROUP BY or an aggregate function
+SELECT y, count(*) FROM t_wide GROUP BY CAST(x AS BIGINT);
Review Comment:
`main` accepts this query and its result is correct. Rejecting it is the
right call, but it is a user-facing change. Could you add it to "Are there any
user-facing changes?" in the description? The upgrade guide already has it.
##########
datafusion/common/src/functional_dependencies.rs:
##########
@@ -480,24 +485,21 @@ pub fn aggregate_functional_dependencies(
{
// Indices into the GROUP BY list for this determinant:
let mut new_source_indices = vec![];
- let mut new_source_field_names = vec![];
- let source_field_names = source_indices
- .iter()
- .map(|&idx| &aggr_input_fields[idx])
- .collect::<Vec<_>>();
-
- for (idx, group_by_expr_name) in
group_by_expr_names.iter().enumerate() {
- // When one of the input determinant expressions matches with
- // the GROUP BY expression, add the index of the GROUP BY
- // expression as a new determinant key:
- if source_field_names.contains(&group_by_expr_name) {
+ let mut new_source_input_indices = vec![];
+ for (idx, input_idx) in group_by_input_indices.iter().enumerate() {
+ // When one of the input determinant columns is a GROUP BY
+ // expression, add the index of the GROUP BY expression as a
new
+ // determinant key:
+ if let Some(input_idx) = input_idx
+ && source_indices.contains(input_idx)
+ {
Review Comment:
A key column that is in the GROUP BY list two times loses its dependency.
This is a regression from `main`, and only the builder / DataFrame API can
reach it (SQL removes aliases from GROUP BY).
`.aggregate(vec![col("id"), col("state"), col("id").alias("k")], ...)` with
`PRIMARY KEY (id)`:
| | Output dependency |
|---|---|
| `main` | `[0] -> [0, 1, 2]` |
| This PR | `[0, 1, 2] -> [0, 1, 2]` |
`id` and `id AS k` both map to input index 0, so the length check below
fails. Results stay correct, but parent plans can no longer prune with `id`.
```suggestion
if let Some(input_idx) = input_idx
&& source_indices.contains(input_idx)
// A key column can be in the GROUP BY list more than one
// time (`x, x AS k`). Count it one time.
&& !new_source_input_indices.contains(&Some(*input_idx))
{
```
A unit test is optional here, because only the builder API reaches this. If
you want one, this fails on this PR and passes with the suggestion:
<details><summary>Optional unit test for <code>plan.rs</code></summary>
```rust
#[test]
fn aggregate_group_by_repeated_key_keeps_key_dependency() -> Result<()> {
let constraints =
Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]);
let source = Arc::new(
LogicalTableSource::new(Arc::new(employee_schema()))
.with_constraints(constraints),
);
// `id` is in the GROUP BY list two times: as itself and as `k`.
let plan = LogicalPlanBuilder::scan("employee_csv", source, None)?
.aggregate(
vec![col("id"), col("state"), col("id").alias("k")],
Vec::<Expr>::new(),
)?
.build()?;
// `id` alone still determines the row.
let deps = plan.schema().functional_dependencies();
assert_eq!(deps.len(), 1);
assert_eq!(deps[0].source_indices, vec![0]);
assert_eq!(deps[0].target_indices, vec![0, 1, 2]);
Ok(())
}
```
</details>
With the suggestion, the full sqllogictest suite and the
`datafusion-common`, `datafusion-expr` and `datafusion-optimizer` unit tests
pass locally.
##########
datafusion/sqllogictest/test_files/functional_dependencies.slt:
##########
@@ -358,6 +358,89 @@ drop table t_null;
statement ok
drop table t_probe;
+
+# 6. Only a column reference carries a dependency. `CAST(x AS INT)` is named
+# `t_cast.x`, like the PRIMARY KEY column, but it does not determine `y`:
+# 1.1 and 1.2 both cast to 1.
+statement ok
+CREATE TABLE t_cast (x DOUBLE, y VARCHAR, PRIMARY KEY (x)) AS VALUES (1.1,
'b'), (1.2, 'a'), (2.1, 'c');
+
+# 6.1 ORDER BY: `y` is the tie-breaker within `CAST(x AS INT) = 1` and is kept.
+query RT
+SELECT x, y FROM t_cast ORDER BY CAST(x AS INT), y;
+----
+1.2 a
+1.1 b
+2.1 c
+
+query TT
+EXPLAIN SELECT x, y FROM t_cast ORDER BY CAST(x AS INT), y;
+----
+logical_plan
+01)Sort: CAST(t_cast.x AS Int32) ASC NULLS LAST, t_cast.y ASC NULLS LAST
+02)--TableScan: t_cast projection=[x, y]
+
+# 6.2 GROUP BY: `y` is not selected but still splits the groups.
+query II rowsort
+SELECT CAST(x AS INT) k, count(*) n FROM t_cast GROUP BY CAST(x AS INT), y;
+----
+1 1
+1 1
+2 1
+
+# 6.3 GROUP BY output: the dependency `k -> y` does not exist, so `y` still
+# orders the rows within `k = 1`.
+query IT
+SELECT CAST(x AS INT) k, y FROM t_cast GROUP BY CAST(x AS INT), y ORDER BY k,
y;
+----
+1 a
+1 b
+2 c
+
+# 6.4 `y` is not determined by the GROUP BY expression, so it can't be
selected.
+query error DataFusion error: Error during planning: Column in SELECT must be
in GROUP BY or an aggregate function
+SELECT y, count(*) FROM t_cast GROUP BY CAST(x AS INT);
+
+statement ok
+drop table t_cast;
+
+# 6.5 Only the key column itself carries the dependency, even when the cast
+# keeps distinct values distinct, like INT to BIGINT. PostgreSQL and DuckDB
+# reject this too.
+statement ok
+CREATE TABLE t_wide (x INT, y VARCHAR, PRIMARY KEY (x)) AS VALUES (1, 'b'),
(2, 'a');
+
+query error DataFusion error: Error during planning: Column in SELECT must be
in GROUP BY or an aggregate function
+SELECT y, count(*) FROM t_wide GROUP BY CAST(x AS BIGINT);
+
+statement ok
+drop table t_wide;
+
+# 6.6 The same with `TRY_CAST` over a `UNIQUE NOT NULL` key: 'bad-a' and
+# 'bad-b' both cast to NULL.
+statement ok
+CREATE TABLE t_try_cast (x VARCHAR NOT NULL UNIQUE, y VARCHAR) AS VALUES
('bad-a', 'b'), ('bad-b', 'a'), ('1', 'c');
+
+query TT
+SELECT x, y FROM t_try_cast ORDER BY TRY_CAST(x AS INT), y;
+----
+1 c
+bad-b a
+bad-a b
+
+query IT rowsort
+SELECT TRY_CAST(x AS INT) k, y FROM t_try_cast GROUP BY TRY_CAST(x AS INT), y;
+----
+1 c
+NULL a
+NULL b
+
+query error DataFusion error: Error during planning: Column in SELECT must be
in GROUP BY or an aggregate function
+SELECT y, count(*) FROM t_try_cast GROUP BY TRY_CAST(x AS INT);
+
+statement ok
+drop table t_try_cast;
Review Comment:
The description says the bug also occurs with a key that comes from an inner
GROUP BY, but no test covers that. The first two queries below return wrong
results on `main`, and `main` accepts the third. All three pass on this PR.
```suggestion
drop table t_try_cast;
# 6.7 A key that comes from an inner GROUP BY: `x` is unique in the subquery,
# but `CAST(x AS INT)` is not.
statement ok
CREATE TABLE t_inner (x DOUBLE, y VARCHAR) AS VALUES (1.1, 'b'), (1.2, 'a'),
(2.1, 'c');
query RT
SELECT x, y FROM (SELECT x, max(y) AS y FROM t_inner GROUP BY x) ORDER BY
CAST(x AS INT), y;
----
1.2 a
1.1 b
2.1 c
query II rowsort
SELECT CAST(x AS INT) k, count(*) n FROM (SELECT x, max(y) AS y FROM t_inner
GROUP BY x) GROUP BY CAST(x AS INT), y;
----
1 1
1 1
2 1
query error DataFusion error: Error during planning: Column in SELECT must
be in GROUP BY or an aggregate function
SELECT y, count(*) FROM (SELECT x, max(y) AS y FROM t_inner GROUP BY x)
GROUP BY CAST(x AS INT);
statement ok
drop table t_inner;
```
--
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]
