namanjain24-sudo commented on code in PR #25527: URL: https://github.com/apache/datafusion/pull/25527#discussion_r4227356038
########## datafusion/substrait/src/logical_plan/grouping_set.rs: ########## @@ -0,0 +1,308 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The column a multi-set aggregate ends with, which Substrait and DataFusion +//! fill differently. +//! +//! Substrait gives an [`AggregateRel`] with more than one grouping set a +//! trailing `i32` holding "the zero-based index of the grouping set that +//! yielded the record" ([Aggregate Operation]). DataFusion ends the same +//! aggregate with `__grouping_id`, which packs a bitmask of the columns the set +//! leaves out together with an ordinal separating repeated sets. Both identify +//! the set a row came from, so each side can be written as a map of the other, +//! which is what the consumer and the producer apply. +//! +//! [`AggregateRel`]: substrait::proto::AggregateRel +//! [Aggregate Operation]: https://substrait.io/relations/logical_relations/#aggregate-operation + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{ + Column, DFSchema, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, +}; +use datafusion::logical_expr::utils::grouping_set_to_exprlist; +use datafusion::logical_expr::{Aggregate, Case, Expr, GroupingSet, lit}; + +/// The preferred name for the grouping set index, which Substrait leaves to +/// the plan's root names. Both the consumer and the producer only use this +/// literal name as the base candidate passed to +/// [`unique_grouping_set_index_name`]; neither hard-codes it as the name a +/// real schema is guaranteed to accept. +const GROUPING_SET_INDEX: &str = "grouping_set_index"; + +/// A name for the grouping set index column that is not already taken in +/// `schema`. +/// +/// `schema` is a real aggregate's output (the consumer) or that aggregate's +/// `DFSchema` sans `__grouping_id` (the producer), so it can legitimately +/// already contain a column named [`GROUPING_SET_INDEX`] - either a user +/// column with that literal name, or (on the producer side specifically) two +/// joined columns that only collide once reduced to this bare, unqualified +/// name. Falling back to the fixed name regardless would make the synthetic +/// column indistinguishable from that real one, or fail schema construction +/// outright with a duplicate-field error. +pub(crate) fn unique_grouping_set_index_name(schema: &DFSchema) -> String { + if schema + .index_of_column_by_name(None, GROUPING_SET_INDEX) + .is_none() + { + return GROUPING_SET_INDEX.to_string(); + } + let mut suffix = 0u32; + loop { + let candidate = format!("{GROUPING_SET_INDEX}_{suffix}"); + if schema.index_of_column_by_name(None, &candidate).is_none() { + return candidate; + } + suffix += 1; + } +} + +/// The `__grouping_id` value DataFusion gives each grouping set, in the order +/// the sets are listed. +/// +/// The value is `(ordinal << group_count) | mask`: a bit is set in `mask` for +/// every grouping column the set leaves out, counting from the last column, and +/// `ordinal` counts the sets before this one holding the same columns. Both +/// parts follow from the set alone, so no two sets share a value. +/// +/// At `group_count == 64` the mask alone already fills every bit of the `u64`, +/// leaving no room for an ordinal: `ordinal << 64` panics on the shift amount +/// even when `ordinal` is `0`, and any `ordinal != 0` cannot be represented at +/// all. The first set at that width is still representable - its ordinal is +/// always `0` - so only a repeated set at exactly 64 columns is rejected. +pub(crate) fn grouping_set_ids( + columns: &[&Expr], + sets: &[Vec<Expr>], +) -> datafusion::common::Result<Vec<u64>> { + let group_count = columns.len(); + if group_count > 64 { + return not_impl_err!( + "Grouping sets with more than 64 columns are not supported" + ); + } + + let mut ids = Vec::with_capacity(sets.len()); + let mut masks = Vec::with_capacity(sets.len()); + for set in sets { + let mut mask = 0u64; + for (position, column) in columns.iter().enumerate() { + if !set.contains(column) { + mask |= 1 << (group_count - 1 - position); + } + } + let ordinal = masks.iter().filter(|seen| **seen == mask).count() as u64; + masks.push(mask); + let id = if group_count == 64 { + if ordinal != 0 { + return not_impl_err!( + "A grouping set with 64 columns cannot be repeated: there are no bits left for a duplicate ordinal" + ); + } + mask + } else { + (ordinal << group_count) | mask Review Comment: Fixed: computed the representable range from the actual bits left above the mask (64 - group_count) and check it before shifting, for every group_count rather than only 64. Verified by reverting: got exactly the predicted collision, [0, 1<<63, 0] for 3x at 63 columns. Added boundary unit tests (last-ok vs first-overflow at 62 and 63 columns) plus a producer round-trip at the ok boundary and a rejection test at the overflow boundary. -- 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]
