Skip to content

Commit a6c003c

Browse files
committed
introduce sql-logic test for multip-group by
1 parent 0c0aab6 commit a6c003c

2 files changed

Lines changed: 107 additions & 70 deletions

File tree

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/dictionary.rs

Lines changed: 43 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
118
#[cfg(test)]
219
mod multi_group_by_dictionary_test {
320
use std::sync::Arc;
@@ -95,7 +112,7 @@ mod multi_group_by_dictionary_test {
95112
mod general_test {
96113
use super::*;
97114

98-
// Covers: all-distinct pairs unique IDs; cross-batch stability (same pair always same ID)
115+
// Covers: all-distinct pairs -> unique IDs; cross-batch stability (same pair always same ID)
99116
#[test]
100117
fn test_distinct_pairs_and_cross_batch_consistency() {
101118
let schema = mk_schema(vec![
@@ -163,54 +180,14 @@ mod multi_group_by_dictionary_test {
163180
vec![0, 1, 0, 2, 2]
164181
);
165182
}
166-
167-
// Covers: 5 columns (3+), Utf8/Utf8View/LargeUtf8 value types, Int32/Int8 key types,
168-
// distinct tuples → unique IDs, repeated tuples → same IDs, null in one column
169-
#[test]
170-
fn test_five_column_grouping() {
171-
let schema = mk_schema(vec![
172-
dict(DataType::Int32, DataType::Utf8),
173-
dict(DataType::Int32, DataType::Utf8View),
174-
dict(DataType::Int32, DataType::LargeUtf8),
175-
dict(DataType::Int8, DataType::Utf8),
176-
dict(DataType::Int32, DataType::Utf8View),
177-
]);
178-
let keys = || int32_keys(vec![Some(0), Some(1), None, Some(0)]);
179-
let rep_keys = || int32_keys(vec![Some(0), Some(1), Some(0), Some(0)]);
180-
let i8k = || int8_keys(vec![Some(0i8), Some(1i8), Some(0i8), Some(0i8)]);
181-
assert_eq!(
182-
intern_all(
183-
schema,
184-
&[
185-
make_dict(keys(), StringArray::from(vec!["a0", "a1"])),
186-
make_dict_view(
187-
rep_keys(),
188-
StringViewArray::from(vec!["b0", "b1"])
189-
),
190-
make_dict_large(
191-
rep_keys(),
192-
LargeStringArray::from(vec!["c0", "c1"])
193-
),
194-
make_dict(i8k(), StringArray::from(vec!["d0", "d1"])),
195-
make_dict_view(
196-
rep_keys(),
197-
StringViewArray::from(vec!["e0", "e1"])
198-
),
199-
],
200-
),
201-
// (a0,b0,c0,d0,e0), (a1,b1,c1,d1,e1), (null,b0,c0,d0,e0), (a0,b0,c0,d0,e0) repeat
202-
vec![0, 1, 2, 0]
203-
);
204-
}
205183
}
206184

207-
// dictionary_encoding: behavior specific to how dictionary arrays encode values
185+
// dictionary_encoding: behavior specific to how dictionary arrays encode values
208186
// key-index shuffling, non-normalized value arrays, bloated dictionaries, mixed key types
209187
mod dictionary_encoding {
210188
use super::*;
211189

212190
// Covers: (foo,bar) ≠ (bar,foo) order sensitivity; same decoded value at different key
213-
// positions cross-batch; non-normalized dict (duplicate entries in values array);
214191
#[test]
215192
fn test_order_sensitivity_and_key_shuffling() {
216193
let schema = mk_schema(vec![
@@ -262,50 +239,46 @@ mod multi_group_by_dictionary_test {
262239
vec![0, 1, 4]
263240
);
264241
}
265-
266-
// Covers: mixed key types (Int8 vs Int32), Utf8 vs LargeUtf8 value types
242+
// Covers: 5 columns , Utf8/Utf8View/LargeUtf8 value types, Int32/Int8 key types,
267243
#[test]
268-
fn test_different_key_types() {
244+
fn test_five_column_grouping() {
269245
let schema = mk_schema(vec![
270-
dict(DataType::Int8, DataType::Utf8),
246+
dict(DataType::Int32, DataType::Utf8),
247+
dict(DataType::Int32, DataType::Utf8View),
271248
dict(DataType::Int32, DataType::LargeUtf8),
249+
dict(DataType::Int8, DataType::Utf8),
250+
dict(DataType::Int32, DataType::Utf8View),
272251
]);
273-
let mut group_values = make_gv(schema);
252+
let keys = || int32_keys(vec![Some(0), Some(1), None, Some(0)]);
253+
let rep_keys = || int32_keys(vec![Some(0), Some(1), Some(0), Some(0)]);
254+
let i8k = || int8_keys(vec![Some(0i8), Some(1i8), Some(0i8), Some(0i8)]);
274255
assert_eq!(
275-
intern(
276-
&mut group_values,
256+
intern_all(
257+
schema,
277258
&[
278-
make_dict(
279-
int8_keys(vec![Some(0i8), Some(1i8)]),
280-
StringArray::from(vec!["alpha", "beta"]),
259+
make_dict(keys(), StringArray::from(vec!["a0", "a1"])),
260+
make_dict_view(
261+
rep_keys(),
262+
StringViewArray::from(vec!["b0", "b1"])
281263
),
282264
make_dict_large(
283-
int32_keys(vec![Some(0), Some(1)]),
284-
LargeStringArray::from(vec!["X", "Y"]),
285-
),
286-
],
287-
),
288-
vec![0, 1]
289-
);
290-
assert_eq!(
291-
intern(
292-
&mut group_values,
293-
&[
294-
make_dict(
295-
int8_keys(vec![Some(1i8), Some(3i8)]),
296-
StringArray::from(vec!["a", "beta", "b", "alpha"]),
265+
rep_keys(),
266+
LargeStringArray::from(vec!["c0", "c1"])
297267
),
298-
make_dict_large(
299-
int32_keys(vec![Some(0), Some(1)]),
300-
LargeStringArray::from(vec!["Y", "X"]),
268+
make_dict(i8k(), StringArray::from(vec!["d0", "d1"])),
269+
make_dict_view(
270+
rep_keys(),
271+
StringViewArray::from(vec!["e0", "e1"])
301272
),
302273
],
303274
),
304-
vec![1, 0]
275+
// (a0,b0,c0,d0,e0), (a1,b1,c1,d1,e1), (null,b0,c0,d0,e0), (a0,b0,c0,d0,e0) repeat
276+
vec![0, 1, 2, 0]
305277
);
306278
}
307279
}
308280

281+
#[cfg(test)]
309282
mod streaming {
310283
use super::*;
311284

datafusion/sqllogictest/test_files/dictionary.slt

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -515,3 +515,67 @@ DROP TABLE dict_hash_10;
515515

516516
statement ok
517517
DROP TABLE dict_hash_src;
518+
519+
# Multi-column GROUP BY with mixed key types. Verifies order sensitivity
520+
statement ok
521+
CREATE TABLE t_multi_dict_grp AS
522+
SELECT
523+
arrow_cast(column1, 'Dictionary(Int8, Utf8)') AS tag_a,
524+
arrow_cast(column2, 'Dictionary(Int32, LargeUtf8)') AS tag_b,
525+
column3 AS val
526+
FROM (VALUES
527+
('alpha', 'X', 10),
528+
('beta', 'Y', 20),
529+
('alpha', 'X', 30),
530+
('beta', 'X', 40),
531+
('alpha', 'Y', 50)
532+
);
533+
534+
query TTI rowsort
535+
SELECT tag_a, tag_b, SUM(val)
536+
FROM t_multi_dict_grp
537+
GROUP BY tag_a, tag_b;
538+
----
539+
alpha X 40
540+
alpha Y 50
541+
beta X 40
542+
beta Y 20
543+
544+
statement ok
545+
DROP TABLE t_multi_dict_grp;
546+
547+
statement ok
548+
CREATE TABLE t_dict_three_tag AS
549+
SELECT
550+
arrow_cast(column1, 'Dictionary(Int8, Utf8View)') AS region,
551+
arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS tier,
552+
arrow_cast(column3, 'Dictionary(Int32, LargeUtf8)') AS env,
553+
column4 AS val
554+
FROM (VALUES
555+
('east', 'gold', 'prod', 1),
556+
('west', 'silver', 'staging', 2),
557+
('east', 'gold', 'prod', 3),
558+
('east', 'gold', 'staging', 4),
559+
(NULL, 'gold', 'prod', 5),
560+
(NULL, 'gold', 'prod', 6)
561+
);
562+
563+
query I
564+
SELECT COUNT(*) FROM (
565+
SELECT region, tier, env FROM t_dict_three_tag GROUP BY region, tier, env
566+
);
567+
----
568+
4
569+
570+
query TTTI rowsort
571+
SELECT region, tier, env, SUM(val)
572+
FROM t_dict_three_tag
573+
GROUP BY region, tier, env;
574+
----
575+
NULL gold prod 11
576+
east gold prod 4
577+
east gold staging 4
578+
west silver staging 2
579+
580+
statement ok
581+
DROP TABLE t_dict_three_tag;

0 commit comments

Comments
 (0)