Fix dictionary handling in AggregationMaskCompiler

trinodb · Apr 3, 2024 · 8681800 · 8681800
1 parent 9d28476
commit 8681800
Showing 1 changed file with 21 additions and 17 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationMaskCompiler.java b/core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationMaskCompiler.java
@@ -162,48 +162,52 @@ public static Constructor<? extends AggregationMaskBuilder> generateAggregationM
 
         // create expression to test if a position is selected
         Variable maskValueBlock = scope.declareVariable(ByteArrayBlock.class, "maskValueBlock");
-        Variable position = scope.declareVariable("position", body, constantInt(0));
-        BytecodeExpression isPositionSelected = testMaskBlock(maskValueBlock, maskBlockMayHaveNull, position);
+        Variable maskValueBlockPosition = scope.declareVariable("maskValueBlockPosition", body, constantInt(0));
+        BytecodeExpression isMaskPositionSelected = testMaskBlock(maskValueBlock, maskBlockMayHaveNull, maskValueBlockPosition);
+
+        Variable pagePosition = scope.declareVariable("pagePosition", body, constantInt(0));
+        BytecodeExpression isPositionSelected = isMaskPositionSelected;
         for (int i = 0; i < nonNullArgs.size(); i++) {
             Variable arg = nonNullArgs.get(i);
             Variable mayHaveNull = nonNullArgMayHaveNulls.get(i);
-            isPositionSelected = and(isPositionSelected, testPositionIsNotNull(arg, mayHaveNull, position));
+            isPositionSelected = and(isPositionSelected, testPositionIsNotNull(arg, mayHaveNull, pagePosition));
         }
 
         // add all positions that pass the tests
         // at this point the mask block can only be a DictionaryBlock, ByteArrayBlock, or null
         Variable selectedPositionsIndex = scope.declareVariable("selectedPositionsIndex", body, constantInt(0));
         Variable rawIds = scope.declareVariable(int[].class, "rawIds");
         Variable rawIdsOffset = scope.declareVariable(int.class, "rawIdsOffset");
-        Variable dictionaryBlockPosition = scope.declareVariable("dictionaryBlockPosition", body, constantInt(0));
         body.append(new IfStatement()
                 .condition(maskBlock.instanceOf(DictionaryBlock.class))
                         .ifTrue(new BytecodeBlock()
                                 .append(maskValueBlock.set(maskBlock.cast(DictionaryBlock.class).invoke("getDictionary", ValueBlock.class).cast(ByteArrayBlock.class)))
                                 .append(rawIds.set(maskBlock.cast(DictionaryBlock.class).invoke("getRawIds", int[].class)))
                                 .append(rawIdsOffset.set(maskBlock.cast(DictionaryBlock.class).invoke("getRawIdsOffset", int.class)))
                                 .append(new ForLoop()
-                                        .initialize(dictionaryBlockPosition.set(constantInt(0)))
-                                        .condition(lessThan(dictionaryBlockPosition, positionCount))
-                                        .update(dictionaryBlockPosition.increment())
+                                        .initialize(pagePosition.set(constantInt(0)))
+                                        .condition(lessThan(pagePosition, positionCount))
+                                        .update(pagePosition.increment())
                                         .body(new BytecodeBlock()
-                                                .append(position.set(rawIds.getElement(add(rawIdsOffset, dictionaryBlockPosition))))
+                                                .append(maskValueBlockPosition.set(rawIds.getElement(add(rawIdsOffset, pagePosition))))
                                                 .append(new IfStatement()
                                                         .condition(isPositionSelected)
                                                         .ifTrue(new BytecodeBlock()
-                                                                .append(selectedPositions.setElement(selectedPositionsIndex, dictionaryBlockPosition))
+                                                                .append(selectedPositions.setElement(selectedPositionsIndex, pagePosition))
                                                                 .append(selectedPositionsIndex.increment()))))))
                         .ifFalse(new BytecodeBlock()
                                 .append(maskValueBlock.set(maskBlock.cast(ByteArrayBlock.class)))
                                 .append(new ForLoop()
-                                        .initialize(position.set(constantInt(0)))
-                                        .condition(lessThan(position, positionCount))
-                                        .update(position.increment())
-                                        .body(new IfStatement()
-                                                .condition(isPositionSelected)
-                                                .ifTrue(new BytecodeBlock()
-                                                        .append(selectedPositions.setElement(selectedPositionsIndex, position))
-                                                        .append(selectedPositionsIndex.increment()))))));
+                                        .initialize(pagePosition.set(constantInt(0)))
+                                        .condition(lessThan(pagePosition, positionCount))
+                                        .update(pagePosition.increment())
+                                        .body(new BytecodeBlock()
+                                                .append(maskValueBlockPosition.set(pagePosition))
+                                                .append(new IfStatement()
+                                                        .condition(isPositionSelected)
+                                                        .ifTrue(new BytecodeBlock()
+                                                                .append(selectedPositions.setElement(selectedPositionsIndex, pagePosition))
+                                                                .append(selectedPositionsIndex.increment())))))));
 
         body.append(invokeStatic(
                 AggregationMask.class,