jayzhan211 commented on code in PR #27113:
URL: https://github.com/apache/datafusion/pull/27113#discussion_r4513489780
##########
datafusion/common/src/functional_dependencies.rs:
##########
@@ -532,7 +574,19 @@ pub fn aggregate_functional_dependencies(
// aggregation:
if !group_by_input_indices.is_empty() {
let count = group_by_input_indices.len();
- let source_indices = (0..count).collect::<Vec<_>>();
+ // Remove extraneous columns (the GROUP BY expressions that the others
+ // determine) from the key, so that the key still exists after a
+ // projection removes them.
+ let source_indices = match get_required_group_by_exprs_indices(
+ aggr_input_schema,
+ group_by_input_indices,
+ ) {
+ Some(mut indices) => {
+ indices.sort_unstable();
+ indices
+ }
+ None => (0..count).collect(),
+ };
Review Comment:
@asolimando thanks for the update! This still rejects a query that returns
correct results on 55.2.0 and main, when the computed expression sits next to a
column the key doesn't determine. The query below returns 3 rows on 55.2.0 and
fails to plan with this PR. My earlier suggestion missed this path: only the
whole GROUP BY key can prove `(i_item_sk, s_date)` is a key here, and
`get_required_group_by_exprs_indices` gives up on any computed expression. The
change below drops a computed expression from that key when the remaining key
columns determine every field it reads. With it, the full sqllogictest suite
passes.
```sql
SELECT i_item_sk, s_date, c
FROM (SELECT i_item_sk, s_date, count(*) AS c FROM item JOIN sales ON
i_item_sk = s_item_sk
GROUP BY i_item_sk, s_date, upper(i_desc)) q
GROUP BY i_item_sk, s_date;
```
<details>
<summary>Change in functional_dependencies.rs</summary>
```diff
@@ -577,16 +577,11 @@ pub fn
aggregate_functional_dependencies_with_computed_inputs(
// Remove extraneous columns (the GROUP BY expressions that the
others
// determine) from the key, so that the key still exists after a
// projection removes them.
- let source_indices = match get_required_group_by_exprs_indices(
+ let source_indices = group_by_key_indices(
aggr_input_schema,
group_by_input_indices,
- ) {
- Some(mut indices) => {
- indices.sort_unstable();
- indices
- }
- None => (0..count).collect(),
- };
+ computed_input_indices,
+ );
@@ -617,6 +612,58 @@ pub fn
aggregate_functional_dependencies_with_computed_inputs(
FunctionalDependencies::new(aggregate_func_dependencies)
}
+/// Returns the GROUP BY expressions that the whole GROUP BY key needs. A
column
+/// that the other columns determine is left out, and so is a deterministic
+/// computed expression whose fields the remaining columns determine.
+fn group_by_key_indices(
+ schema: &DFSchema,
+ group_by_input_indices: &[Option<usize>],
+ computed_input_indices: &[Option<Vec<usize>>],
+) -> Vec<usize> {
+ let all = || (0..group_by_input_indices.len()).collect();
+ let (column_positions, column_inputs): (Vec<usize>, Vec<Option<usize>>)
=
+ group_by_input_indices
+ .iter()
+ .enumerate()
+ .filter(|(_, input_idx)| input_idx.is_some())
+ .map(|(idx, input_idx)| (idx, *input_idx))
+ .unzip();
+ let Some(required) = get_required_group_by_exprs_indices(schema,
&column_inputs)
+ else {
+ return all();
+ };
+ let key_inputs = required
+ .iter()
+ .map(|idx| column_inputs[*idx])
+ .collect::<Vec<_>>();
+ let determined =
+ get_target_functional_dependencies(schema,
&key_inputs).unwrap_or_default();
+ let mut key = required
+ .iter()
+ .map(|idx| column_positions[*idx])
+ .collect::<Vec<_>>();
+ for (idx, (input_idx, read)) in group_by_input_indices
+ .iter()
+ .zip(computed_input_indices)
+ .enumerate()
+ {
+ let is_determined = read.as_ref().is_some_and(|read| {
+ read.iter().all(|field| {
+ determined.contains(field) ||
key_inputs.contains(&Some(*field))
+ })
+ });
+ if input_idx.is_none() && !is_determined {
+ key.push(idx);
+ }
+ }
+ // A key with no columns is left as the whole GROUP BY list.
+ if key.is_empty() {
+ return all();
+ }
+ key.sort_unstable();
+ key
+}
+
```
</details>
<details>
<summary>Tests to add after 8.5 (the first fails on this PR; the second
fails without the empty-key fallback, where count(*) returns no rows)</summary>
```
# `upper(i_desc)` reads `i_desc`, which `i_item_sk` determines, and `s_date
+ 1`
# reads `s_date`: both are removed from the key of the inner aggregate.
query III
SELECT i_item_sk, s_date, c
FROM (SELECT i_item_sk, s_date, count(*) AS c FROM item JOIN sales ON
i_item_sk = s_item_sk
GROUP BY i_item_sk, s_date, upper(i_desc), s_date + 1) q
GROUP BY i_item_sk, s_date ORDER BY i_item_sk, s_date;
----
1 100 1
1 101 1
2 100 1
# A GROUP BY that reads no column keeps the whole GROUP BY list as its key.
query I
SELECT count(*) FROM orders
JOIN (SELECT 1 AS k FROM sales WHERE false GROUP BY 'x') g ON o_item_sk =
g.k;
----
0
```
</details>
--
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]
