Repository: apache/flink
File: flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PruneAggregateCallRule.java (Java)

Code before:
    }

    private int[] getUnrefAggCallIndices(ImmutableBitSet inputRefs, Aggregate agg) {
        int groupCount = agg.getGroupCount();
        return IntStream.range(0, agg.getAggCallList().size())
                .filter(index -> !inputRefs.get(groupCount + index))
                .toArray();
    }

    @Override
    public void onMatch(RelOptRuleCall call) {
        T relOnAgg = call.rel(0);
        Aggregate agg = call.rel(1);
        ImmutableBitSet inputRefs = getInputRefs(relOnAgg);
        int[] unrefAggCallIndices = getUnrefAggCallIndices(inputRefs, agg);
        Preconditions.checkArgument(unrefAggCallIndices.length > 0, "requirement failed");

        List<AggregateCall> newAggCalls = new ArrayList<>(agg.getAggCallList());
        // remove unreferenced AggCall from original aggCalls
        Arrays.stream(unrefAggCallIndices)
                .boxed()
                .sorted(Comparator.reverseOrder())
                .forEach(index -> newAggCalls.remove((int) index));

        if (newAggCalls.isEmpty() && agg.getGroupCount() == 0) {
            // at least output one column
            newAggCalls.add(agg.getAggCallList().get(0));
            unrefAggCallIndices =
                    Arrays.copyOfRange(unrefAggCallIndices, 1, unrefAggCallIndices.length);
        }

        Aggregate newAgg =
                agg.copy(
                        agg.getTraitSet(),
                        agg.getInput(),
                        agg.getGroupSet(),
                        List.of(agg.getGroupSet()),
                        newAggCalls);

        int newFieldIndex = 0;
        // map old agg output index to new agg output index
        Map<Integer, Integer> mapOldToNew = new HashMap<>();
        int fieldCountOfOldAgg = agg.getRowType().getFieldCount();
        List<Integer> unrefAggCallOutputIndices =
                Arrays.stream(unrefAggCallIndices)
                        .mapToObj(i -> i + agg.getGroupCount())
                        .collect(Collectors.toList());
        for (int i = 0; i < fieldCountOfOldAgg; i++) {
            if (!unrefAggCallOutputIndices.contains(i)) {
                mapOldToNew.put(i, newFieldIndex);
                newFieldIndex++;

Review comment (suggestion):
comments: newAggCalls.remove((int) index) , we need this int cast here. because it will remove the value instead of removing index when doesn't have int cast



Code after:
    }

    private int[] getUnrefAggCallIndices(ImmutableBitSet inputRefs, Aggregate agg) {
        int groupCount = agg.getGroupCount();
        return IntStream.range(0, agg.getAggCallList().size())
                .filter(index -> !inputRefs.get(groupCount + index))
                .toArray();
    }

    @Override
    public void onMatch(RelOptRuleCall call) {
        T relOnAgg = call.rel(0);
        Aggregate agg = call.rel(1);
        ImmutableBitSet inputRefs = getInputRefs(relOnAgg);
        int[] unrefAggCallIndices = getUnrefAggCallIndices(inputRefs, agg);
        Preconditions.checkArgument(unrefAggCallIndices.length > 0, "requirement failed");

        List<AggregateCall> newAggCalls = new ArrayList<>(agg.getAggCallList());
        // remove unreferenced AggCall from original aggCalls
        Arrays.stream(unrefAggCallIndices)
                .boxed()
                .sorted(Comparator.reverseOrder())
                // we need this int cast here. because it will remove the value instead of removing
                // index when doesn't have int cast
                .forEach(index -> newAggCalls.remove((int) index));

        if (newAggCalls.isEmpty() && agg.getGroupCount() == 0) {
            // at least output one column
            newAggCalls.add(agg.getAggCallList().get(0));
            unrefAggCallIndices =
                    Arrays.copyOfRange(unrefAggCallIndices, 1, unrefAggCallIndices.length);
        }

        Aggregate newAgg =
                agg.copy(
                        agg.getTraitSet(),
                        agg.getInput(),
                        agg.getGroupSet(),
                        List.of(agg.getGroupSet()),
                        newAggCalls);

        int newFieldIndex = 0;
        // map old agg output index to new agg output index
        Map<Integer, Integer> mapOldToNew = new HashMap<>();
        int fieldCountOfOldAgg = agg.getRowType().getFieldCount();
        List<Integer> unrefAggCallOutputIndices =
                Arrays.stream(unrefAggCallIndices)
                        .mapToObj(i -> i + agg.getGroupCount())
                        .collect(Collectors.toList());
        for (int i = 0; i < fieldCountOfOldAgg; i++) {
            if (!unrefAggCallOutputIndices.contains(i)) {