From 84346c19efb5f01f0e627968e7cedad191ca9bd6 Mon Sep 17 00:00:00 2001 From: Aleksandr Efimov Date: Thu, 8 Oct 2026 14:03:11 +0300 Subject: [PATCH] fix: avoid shifting zero grouping ordinals by 64 A unique 64-column grouping set fits the UInt64 representation, but packing its zero duplicate ordinal shifts by the full width. Use the semantic mask directly for ordinal zero and preserve the existing capacity checks for nonzero ordinals. Add unit boundary tests and SQL coverage for both aggregate kernels, including real NULL keys, duplicate sets and capacity refusals. --- .../physical-plan/src/aggregates/mod.rs | 57 +++++- .../sqllogictest/test_files/grouping_wide.slt | 170 ++++++++++++++++++ 2 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 datafusion/sqllogictest/test_files/grouping_wide.slt diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 6b56bd308b72b..f37c9de43f59e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3577,7 +3577,13 @@ pub(crate) fn group_id_array( let semantic_id = group .iter() .fold(0u64, |acc, &is_null| (acc << 1) | u64::from(is_null)); - let full_id = semantic_id | ((ordinal as u64) << n); + // A zero ordinal needs no high bits. Avoid shifting it by 64 when all + // UInt64 bits belong to the semantic mask. + let full_id = if ordinal == 0 { + semantic_id + } else { + semantic_id | ((ordinal as u64) << n) + }; if total_bits <= 8 { Ok(Arc::new(UInt8Array::from(vec![full_id as u8; num_rows]))) } else if total_bits <= 16 { @@ -3745,6 +3751,55 @@ mod tests { use futures::{FutureExt, Stream, StreamExt}; use insta::{allow_duplicates, assert_snapshot}; + #[test] + fn wide_group_id_64_zero_ordinal() { + let mut first = [false; 64]; + first[0] = true; + let mut last = [false; 64]; + last[63] = true; + for (group, expected) in [ + ([false; 64], 0), + ([true; 64], u64::MAX), + (first, 9_223_372_036_854_775_808), + (last, 1), + ] { + let array = group_id_array(&group, 0, 0, 2).unwrap(); + assert_eq!(array.data_type(), &DataType::UInt64); + let actual = array.as_any().downcast_ref::().unwrap(); + assert_eq!(actual.values().as_ref(), &[expected, expected]); + } + } + + #[test] + fn wide_group_id_64_empty_rows() { + let array = group_id_array(&[true; 64], 0, 0, 0).unwrap(); + assert_eq!(array.data_type(), &DataType::UInt64); + assert_eq!(array.len(), 0); + } + + #[test] + fn wide_group_id_63_duplicate_bits() { + // 63 omitted-key bits plus the first duplicate fit exactly in UInt64. + let array = group_id_array(&[true; 63], 1, 1, 2).unwrap(); + assert_eq!(array.data_type(), &DataType::UInt64); + let actual = array.as_any().downcast_ref::().unwrap(); + assert_eq!(actual.values().as_ref(), &[u64::MAX, u64::MAX]); + } + + #[test] + fn wide_group_id_capacity_refusals() { + for (width, ordinal, message) in [ + (63, 2, "require 65 bits"), + (64, 1, "require 65 bits"), + (65, 0, "more than 64 columns"), + ] { + let error = + group_id_array(&vec![false; width], ordinal, ordinal, 1).unwrap_err(); + assert!(matches!(&error, DataFusionError::NotImplemented(_))); + assert!(error.to_string().contains(message), "{error}"); + } + } + #[cfg(feature = "proto")] #[test] fn split_human_display_alias_ignores_mismatched_alias() { diff --git a/datafusion/sqllogictest/test_files/grouping_wide.slt b/datafusion/sqllogictest/test_files/grouping_wide.slt new file mode 100644 index 0000000000000..b4ee5a7f1072e --- /dev/null +++ b/datafusion/sqllogictest/test_files/grouping_wide.slt @@ -0,0 +1,170 @@ +# 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. + +# 64 semantic bits fit UInt64 when no duplicate-ordinal bits are needed. +# Real NULL keys must remain distinct from rolled-up NULLs. +statement ok +CREATE TABLE wide_keys ( + c0 INTEGER, + c1 INTEGER, + c2 INTEGER, + c3 INTEGER, + c4 INTEGER, + c5 INTEGER, + c6 INTEGER, + c7 INTEGER, + c8 INTEGER, + c9 INTEGER, + c10 INTEGER, + c11 INTEGER, + c12 INTEGER, + c13 INTEGER, + c14 INTEGER, + c15 INTEGER, + c16 INTEGER, + c17 INTEGER, + c18 INTEGER, + c19 INTEGER, + c20 INTEGER, + c21 INTEGER, + c22 INTEGER, + c23 INTEGER, + c24 INTEGER, + c25 INTEGER, + c26 INTEGER, + c27 INTEGER, + c28 INTEGER, + c29 INTEGER, + c30 INTEGER, + c31 INTEGER, + c32 INTEGER, + c33 INTEGER, + c34 INTEGER, + c35 INTEGER, + c36 INTEGER, + c37 INTEGER, + c38 INTEGER, + c39 INTEGER, + c40 INTEGER, + c41 INTEGER, + c42 INTEGER, + c43 INTEGER, + c44 INTEGER, + c45 INTEGER, + c46 INTEGER, + c47 INTEGER, + c48 INTEGER, + c49 INTEGER, + c50 INTEGER, + c51 INTEGER, + c52 INTEGER, + c53 INTEGER, + c54 INTEGER, + c55 INTEGER, + c56 INTEGER, + c57 INTEGER, + c58 INTEGER, + c59 INTEGER, + c60 INTEGER, + c61 INTEGER, + c62 INTEGER, + c63 INTEGER, + c64 INTEGER +); + +statement ok +INSERT INTO wide_keys VALUES +(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65), +(NULL, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, NULL, 165); + +statement ok +SET datafusion.execution.enable_migration_aggregate = false; + +query IIIII rowsort +SELECT c0, c63, grouping(c0) AS g0, grouping(c63) AS g63, count(*) AS n +FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63), (c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63), (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62), ()); +---- +1 64 0 0 1 +1 NULL 0 1 1 +NULL 64 1 0 1 +NULL NULL 0 0 1 +NULL NULL 0 1 1 +NULL NULL 1 0 1 +NULL NULL 1 1 2 + +# 63 semantic bits plus one duplicate-ordinal bit still fit UInt64. +query IIIII rowsort +SELECT c0, c62, grouping(c0), grouping(c62), count(*) +FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62), (), (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62), ()); +---- +1 63 0 0 1 +1 63 0 0 1 +NULL 163 0 0 1 +NULL 163 0 0 1 +NULL NULL 1 1 2 +NULL NULL 1 1 2 + +# A duplicate at 64 keys still needs 65 bits and must return an error. +statement error require 65 bits, which exceeds 64 +SELECT count(*) FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63), (), (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63)); + +# More than 64 distinct keys remain outside the native representation. +statement error Grouping sets with more than 64 columns are not supported +SELECT count(*) FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63, c64), ()); + +statement ok +SET datafusion.execution.enable_migration_aggregate = true; + +query IIIII rowsort +SELECT c0, c63, grouping(c0) AS g0, grouping(c63) AS g63, count(*) AS n +FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63), (c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63), (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62), ()); +---- +1 64 0 0 1 +1 NULL 0 1 1 +NULL 64 1 0 1 +NULL NULL 0 0 1 +NULL NULL 0 1 1 +NULL NULL 1 0 1 +NULL NULL 1 1 2 + +# 63 semantic bits plus one duplicate-ordinal bit still fit UInt64. +query IIIII rowsort +SELECT c0, c62, grouping(c0), grouping(c62), count(*) +FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62), (), (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62), ()); +---- +1 63 0 0 1 +1 63 0 0 1 +NULL 163 0 0 1 +NULL 163 0 0 1 +NULL NULL 1 1 2 +NULL NULL 1 1 2 + +# A duplicate at 64 keys still needs 65 bits and must return an error. +statement error require 65 bits, which exceeds 64 +SELECT count(*) FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63), (), (c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63)); + +# More than 64 distinct keys remain outside the native representation. +statement error Grouping sets with more than 64 columns are not supported +SELECT count(*) FROM wide_keys +GROUP BY GROUPING SETS ((c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c20, c21, c22, c23, c24, c25, c26, c27, c28, c29, c30, c31, c32, c33, c34, c35, c36, c37, c38, c39, c40, c41, c42, c43, c44, c45, c46, c47, c48, c49, c50, c51, c52, c53, c54, c55, c56, c57, c58, c59, c60, c61, c62, c63, c64), ());