Skip to content

Commit 30d35e3

Browse files
gforsythgforsyth
committed
fix(grouping): collect multiple grouping keys in single "groupings"
Also updated our selection routine to properly handle the flattened index count for projecting out of these groupbys.
1 parent d2c96a3 commit 30d35e3

11 files changed

Lines changed: 57 additions & 88 deletions

File tree

‎ibis_substrait/compiler/translate.py‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -856,7 +856,27 @@ def selection(
856856
for relname, rel in rels:
857857
if relname == "aggregate" and rel.measures:
858858
mapping_counter = itertools.count(
859-
len(rel.measures) + len(rel.groupings)
859+
len(rel.measures)
860+
# Individual groups can have multiple grouping expressions
861+
# and we need to count all of these to properly index projections
862+
# e.g. for the following query
863+
# SELECT
864+
# t0.b,
865+
# t0.sum
866+
# FROM (
867+
# SELECT
868+
# t1.a AS a,
869+
# t1.b AS b,
870+
# SUM(t1.c) AS sum
871+
# FROM t AS t1
872+
# GROUP BY
873+
# t1.a,
874+
# t1.b
875+
# ) AS t0
876+
#
877+
# the two grouping keys (t1.a, t1.b) will be grouping
878+
# expressions in the first (and only) group.
879+
+ sum(len(group.grouping_expressions) for group in rel.groupings)
860880
)
861881
break
862882
elif output_mapping := rel.common.emit.output_mapping:
@@ -1092,9 +1112,10 @@ def aggregation(
10921112
input=input,
10931113
groupings=[
10941114
stalg.AggregateRel.Grouping(
1095-
grouping_expressions=[translate(by, compiler=compiler, **kwargs)]
1115+
grouping_expressions=[
1116+
translate(by, compiler=compiler, **kwargs) for by in op.by
1117+
]
10961118
)
1097-
for by in op.by
10981119
],
10991120
measures=[
11001121
stalg.AggregateRel.Measure(

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h01/tpc_h01.json‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -243,11 +243,7 @@
243243
},
244244
"rootReference": {}
245245
}
246-
}
247-
]
248-
},
249-
{
250-
"groupingExpressions": [
246+
},
251247
{
252248
"selection": {
253249
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h03/tpc_h03.json‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -567,11 +567,7 @@
567567
},
568568
"rootReference": {}
569569
}
570-
}
571-
]
572-
},
573-
{
574-
"groupingExpressions": [
570+
},
575571
{
576572
"selection": {
577573
"directReference": {
@@ -581,11 +577,7 @@
581577
},
582578
"rootReference": {}
583579
}
584-
}
585-
]
586-
},
587-
{
588-
"groupingExpressions": [
580+
},
589581
{
590582
"selection": {
591583
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h06/tpc_h06.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,9 @@
370370
}
371371
}
372372
},
373+
"groupings": [
374+
{}
375+
],
373376
"measures": [
374377
{
375378
"measure": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1128,11 +1128,7 @@
11281128
},
11291129
"rootReference": {}
11301130
}
1131-
}
1132-
]
1133-
},
1134-
{
1135-
"groupingExpressions": [
1131+
},
11361132
{
11371133
"selection": {
11381134
"directReference": {
@@ -1142,11 +1138,7 @@
11421138
},
11431139
"rootReference": {}
11441140
}
1145-
}
1146-
]
1147-
},
1148-
{
1149-
"groupingExpressions": [
1141+
},
11501142
{
11511143
"selection": {
11521144
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h091/tpc_h091.json‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1037,11 +1037,7 @@
10371037
},
10381038
"rootReference": {}
10391039
}
1040-
}
1041-
]
1042-
},
1043-
{
1044-
"groupingExpressions": [
1040+
},
10451041
{
10461042
"selection": {
10471043
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h092/tpc_h092.json‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1025,11 +1025,7 @@
10251025
},
10261026
"rootReference": {}
10271027
}
1028-
}
1029-
]
1030-
},
1031-
{
1032-
"groupingExpressions": [
1028+
},
10331029
{
10341030
"selection": {
10351031
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h10/tpc_h10.json‎

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -651,11 +651,7 @@
651651
},
652652
"rootReference": {}
653653
}
654-
}
655-
]
656-
},
657-
{
658-
"groupingExpressions": [
654+
},
659655
{
660656
"selection": {
661657
"directReference": {
@@ -665,11 +661,7 @@
665661
},
666662
"rootReference": {}
667663
}
668-
}
669-
]
670-
},
671-
{
672-
"groupingExpressions": [
664+
},
673665
{
674666
"selection": {
675667
"directReference": {
@@ -679,11 +671,7 @@
679671
},
680672
"rootReference": {}
681673
}
682-
}
683-
]
684-
},
685-
{
686-
"groupingExpressions": [
674+
},
687675
{
688676
"selection": {
689677
"directReference": {
@@ -693,11 +681,7 @@
693681
},
694682
"rootReference": {}
695683
}
696-
}
697-
]
698-
},
699-
{
700-
"groupingExpressions": [
684+
},
701685
{
702686
"selection": {
703687
"directReference": {
@@ -707,11 +691,7 @@
707691
},
708692
"rootReference": {}
709693
}
710-
}
711-
]
712-
},
713-
{
714-
"groupingExpressions": [
694+
},
715695
{
716696
"selection": {
717697
"directReference": {
@@ -721,11 +701,7 @@
721701
},
722702
"rootReference": {}
723703
}
724-
}
725-
]
726-
},
727-
{
728-
"groupingExpressions": [
704+
},
729705
{
730706
"selection": {
731707
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h18/tpc_h18.json‎

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -427,23 +427,15 @@
427427
},
428428
"rootReference": {}
429429
}
430-
}
431-
]
432-
},
433-
{
434-
"groupingExpressions": [
430+
},
435431
{
436432
"selection": {
437433
"directReference": {
438434
"structField": {}
439435
},
440436
"rootReference": {}
441437
}
442-
}
443-
]
444-
},
445-
{
446-
"groupingExpressions": [
438+
},
447439
{
448440
"selection": {
449441
"directReference": {
@@ -453,11 +445,7 @@
453445
},
454446
"rootReference": {}
455447
}
456-
}
457-
]
458-
},
459-
{
460-
"groupingExpressions": [
448+
},
461449
{
462450
"selection": {
463451
"directReference": {
@@ -467,11 +455,7 @@
467455
},
468456
"rootReference": {}
469457
}
470-
}
471-
]
472-
},
473-
{
474-
"groupingExpressions": [
458+
},
475459
{
476460
"selection": {
477461
"directReference": {

‎ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h19/tpc_h19.json‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1322,6 +1322,9 @@
13221322
}
13231323
}
13241324
},
1325+
"groupings": [
1326+
{}
1327+
],
13251328
"measures": [
13261329
{
13271330
"measure": {

0 commit comments

Comments
 (0)