Skip to content

Commit

Permalink
Fix unparsing of Aggregate with grouping sets
Browse files Browse the repository at this point in the history
  • Loading branch information
eejbyfeldt committed Oct 6, 2024
1 parent 920b384 commit 1f61ddf
Showing 1 changed file with 13 additions and 4 deletions.
17 changes: 13 additions & 4 deletions datafusion/sql/src/unparser/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
// specific language governing permissions and limitations
// under the License.

use std::cmp::Ordering;

use datafusion_common::{
internal_err,
tree_node::{Transformed, TreeNode},
Expand Down Expand Up @@ -166,10 +168,17 @@ fn find_agg_expr<'a>(agg: &'a Aggregate, column: &Column) -> Result<Option<&'a E
if matches!(agg.group_expr.as_slice(), [Expr::GroupingSet(_)]) {
// For grouping set expr, we must operate by expression list from the grouping set
let grouping_expr = grouping_set_to_exprlist(agg.group_expr.as_slice())?;
Ok(grouping_expr
.into_iter()
.chain(agg.aggr_expr.iter())
.nth(index))
match index.cmp(&grouping_expr.len()) {
Ordering::Less => Ok(grouping_expr.into_iter().nth(index)),
Ordering::Equal => {
internal_err!(
"Tried to unproject column refereing to internal grouping id"
)
}
Ordering::Greater => {
Ok(agg.aggr_expr.get(index - grouping_expr.len() - 1))
}
}
} else {
Ok(agg.group_expr.iter().chain(agg.aggr_expr.iter()).nth(index))
}
Expand Down

0 comments on commit 1f61ddf

Please sign in to comment.