Change GroupedAccumulatorState groupId from long to int

Author: dain

core/trino-main/src/main/java/io/trino/operator/aggregation/AccumulatorCompiler.java

@@ -356,7 +356,7 @@ private static void generateGetEstimatedSize(ClassDefinition definition, List<Fi
 
     private static void generateSetGroupCount(ClassDefinition definition, List<FieldDefinition> stateFields)
     {
-        Parameter groupCount = arg("groupCount", long.class);
+        Parameter groupCount = arg("groupCount", int.class);
 
         MethodDefinition method = definition.declareMethod(a(PUBLIC), "setGroupCount", type(void.class), groupCount);
         BytecodeBlock body = method.getBody();
@@ -781,7 +781,7 @@ private static void generateAddIntermediateAsCombine(
         for (FieldDefinition stateField : stateFields) {
             if (grouped) {
                 Variable groupIds = scope.getVariable("groupIds");
-                loopBody.append(thisVariable.getField(stateField).invoke("setGroupId", void.class, groupIds.getElement(position).cast(long.class)));
+                loopBody.append(thisVariable.getField(stateField).invoke("setGroupId", void.class, groupIds.getElement(position)));
             }
             loopBody.append(thisVariable.getField(stateField));
         }
@@ -805,7 +805,7 @@ private static void generateSetGroupIdFromGroupIds(Scope scope, List<FieldDefini
         Variable groupIds = scope.getVariable("groupIds");
         for (FieldDefinition stateField : stateFields) {
             BytecodeExpression state = scope.getThis().getField(stateField);
-            block.append(state.invoke("setGroupId", void.class, groupIds.getElement(position).cast(long.class)));
+            block.append(state.invoke("setGroupId", void.class, groupIds.getElement(position)));
         }
     }
 
@@ -875,14 +875,14 @@ private static void generateGroupedEvaluateIntermediate(ClassDefinition definiti
             BytecodeExpression stateSerializer = thisVariable.getField(getOnlyElement(stateFieldAndDescriptors).getStateSerializerField());
             BytecodeExpression state = thisVariable.getField(getOnlyElement(stateFieldAndDescriptors).getStateField());
 
-            body.append(state.invoke("setGroupId", void.class, groupId.cast(long.class)))
+            body.append(state.invoke("setGroupId", void.class, groupId))
                     .append(stateSerializer.invoke("serialize", void.class, state.cast(AccumulatorState.class), out))
                     .ret();
         }
         else {
             for (StateFieldAndDescriptor stateFieldAndDescriptor : stateFieldAndDescriptors) {
                 BytecodeExpression state = thisVariable.getField(stateFieldAndDescriptor.getStateField());
-                body.append(state.invoke("setGroupId", void.class, groupId.cast(long.class)));
+                body.append(state.invoke("setGroupId", void.class, groupId));
             }
 
             generateSerializeState(definition, stateFieldAndDescriptors, out, thisVariable, body);
@@ -966,7 +966,7 @@ private static void generateGroupedEvaluateFinal(
 
         for (FieldDefinition stateField : stateFields) {
             BytecodeExpression state = thisVariable.getField(stateField);
-            body.append(state.invoke("setGroupId", void.class, groupId.cast(long.class)));
+            body.append(state.invoke("setGroupId", void.class, groupId));
             states.add(state);
         }
 

core/trino-main/src/main/java/io/trino/operator/aggregation/AggregationLoopBuilder.java

@@ -194,7 +194,7 @@ private static MethodDefinition buildCoreLoop(
             Variable groupId = scope.declareVariable(int.class, "groupId");
             invokeFunction.append(groupId.set(aggregationParameters.groupIds().get().getElement(position)));
             for (Parameter stateParameter : aggregationParameters.states()) {
-                invokeFunction.append(stateParameter.cast(GroupedAccumulatorState.class).invoke("setGroupId", void.class, groupId.cast(long.class)));
+                invokeFunction.append(stateParameter.cast(GroupedAccumulatorState.class).invoke("setGroupId", void.class, groupId));
             }
         }
         invokeFunction.append(invoke(binder.bind(function), "input", aggregationArguments.build()));

core/trino-main/src/main/java/io/trino/operator/aggregation/DistinctAccumulatorFactory.java

@@ -202,7 +202,7 @@ public long getEstimatedSize()
         }
 
         @Override
-        public void setGroupCount(long groupCount)
+        public void setGroupCount(int groupCount)
         {
             accumulator.setGroupCount(groupCount);
         }

core/trino-main/src/main/java/io/trino/operator/aggregation/GroupedAccumulator.java

@@ -21,7 +21,7 @@ public interface GroupedAccumulator
 {
     long getEstimatedSize();
 
-    void setGroupCount(long groupCount);
+    void setGroupCount(int groupCount);
 
     void addInput(int[] groupIds, Page page, AggregationMask mask);
 

core/trino-main/src/main/java/io/trino/operator/aggregation/GroupedMapAggregationState.java

@@ -20,8 +20,6 @@
 
 import java.lang.invoke.MethodHandle;
 
-import static java.lang.Math.toIntExact;
-
 public class GroupedMapAggregationState
         extends AbstractMapAggregationState
         implements GroupedAccumulatorState
@@ -53,15 +51,15 @@ public GroupedMapAggregationState(
     }
 
     @Override
-    public void setGroupId(long groupId)
+    public void setGroupId(int groupId)
     {
-        this.groupId = toIntExact(groupId);
+        this.groupId = groupId;
     }
 
     @Override
-    public void ensureCapacity(long size)
+    public void ensureCapacity(int size)
     {
-        setMaxGroupId(toIntExact(size));
+        setMaxGroupId(size);
     }
 
     @Override

core/trino-main/src/main/java/io/trino/operator/aggregation/LongApproximateMostFrequentStateFactory.java

@@ -84,7 +84,7 @@ public void set(ApproximateMostFrequentHistogram<Long> histogram)
         }
 
         @Override
-        public void ensureCapacity(long size)
+        public void ensureCapacity(int size)
         {
             histograms.ensureCapacity(size);
         }

core/trino-main/src/main/java/io/trino/operator/aggregation/NumericHistogramStateFactory.java

@@ -42,7 +42,7 @@ public static class GroupedState
         private long size;
 
         @Override
-        public void ensureCapacity(long size)
+        public void ensureCapacity(int size)
         {
             histograms.ensureCapacity(size);
         }

core/trino-main/src/main/java/io/trino/operator/aggregation/OrderedAccumulatorFactory.java

@@ -207,7 +207,7 @@ public long getEstimatedSize()
         }
 
         @Override
-        public void setGroupCount(long groupCount)
+        public void setGroupCount(int groupCount)
         {
             this.groupCount = max(this.groupCount, groupCount);
             accumulator.setGroupCount(groupCount);

core/trino-main/src/main/java/io/trino/operator/aggregation/StringApproximateMostFrequentStateFactory.java

@@ -85,7 +85,7 @@ public void set(ApproximateMostFrequentHistogram<Slice> histogram)
         }
 
         @Override
-        public void ensureCapacity(long size)
+        public void ensureCapacity(int size)
         {
             histograms.ensureCapacity(size);
         }

core/trino-main/src/main/java/io/trino/operator/aggregation/arrayagg/GroupArrayAggregationState.java

@@ -25,7 +25,6 @@
 import static io.airlift.slice.SizeOf.instanceSize;
 import static io.airlift.slice.SizeOf.sizeOf;
 import static java.lang.Math.clamp;
-import static java.lang.Math.toIntExact;
 
 public final class GroupArrayAggregationState
         extends AbstractGroupedAccumulatorState
@@ -55,10 +54,10 @@ public long getEstimatedSize()
     }
 
     @Override
-    public void ensureCapacity(long maxGroupId)
+    public void ensureCapacity(int maxGroupId)
     {
         checkArgument(maxGroupId + 1 < MAX_ARRAY_SIZE, "Maximum array size exceeded");
-        int requiredSize = toIntExact(maxGroupId + 1);
+        int requiredSize = maxGroupId + 1;
         if (requiredSize > groupHeadPositions.length) {
             int newSize = clamp(requiredSize * 2L, 1024, MAX_ARRAY_SIZE);
             int oldSize = groupHeadPositions.length;