POC sum_with_keep(col, 'type') within group

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/exp/agg/Accumulators.java

@@ -78,6 +78,8 @@ private static <Row> Supplier<Accumulator<Row>> accumulatorFunctionFactory(
                 return avgFactory(call, hnd);
             case "SUM":
                 return sumFactory(call, hnd);
+            case "SUM_WITH_KEEP":
+                return sumWithKeepFactory(call, ctx);
             case "$SUM0":
                 return sumEmptyIsZeroFactory(call, hnd);
             case "MIN":
@@ -213,6 +215,34 @@ private static <Row> Supplier<Accumulator<Row>> sumEmptyIsZeroFactory(AggregateC
         }
     }
 
+    /** */
+    private static <Row> Supplier<Accumulator<Row>> sumWithKeepFactory(
+        AggregateCall call,
+        ExecutionContext<Row> ctx
+    ) {
+        assert call.getCollation() != null && !call.getCollation().getFieldCollations().isEmpty();
+
+        RowHandler<Row> hnd = ctx.rowHandler();
+        Comparator<Row> cmp = ctx.expressionFactory().comparator(call.getCollation());
+
+        switch (call.type.getSqlTypeName()) {
+            case BIGINT:
+            case DECIMAL:
+                return () -> new SumWithKeep<>(call, hnd, cmp, DECIMAL);
+
+            case DOUBLE:
+            case REAL:
+            case FLOAT:
+                return () -> new SumWithKeep<>(call, hnd, cmp, DOUBLE);
+
+            case TINYINT:
+            case SMALLINT:
+            case INTEGER:
+            default:
+                return () -> new SumWithKeep<>(call, hnd, cmp, BIGINT);
+        }
+    }
+
     /** */
     private static <Row> Supplier<Accumulator<Row>> minFactory(AggregateCall call, RowHandler<Row> hnd) {
         switch (call.type.getSqlTypeName()) {
@@ -696,6 +726,131 @@ public Sum(AggregateCall aggCall, Accumulator<Row> acc, RowHandler<Row> hnd) {
         }
     }
 
+    /** SUM(value) over rows tied at the first or last ordering key. */
+    private static class SumWithKeep<Row> extends AbstractAccumulator<Row> implements StoringAccumulator {
+        /** Type used for intermediate sum. */
+        private final SqlTypeName sumType;
+
+        /** Comparator for WITHIN GROUP collation. */
+        private final transient Comparator<Row> cmp;
+
+        /** Row having the selected ordering key. */
+        private Row bestRow;
+
+        /** Intermediate sum. */
+        private Object sum;
+
+        /** Whether at least one row was seen. */
+        private boolean touched;
+
+        /** Whether the first ordering key should be kept. */
+        private boolean first;
+
+        /** */
+        SumWithKeep(AggregateCall aggCall, RowHandler<Row> hnd, Comparator<Row> cmp, SqlTypeName sumType) {
+            super(aggCall, hnd);
+
+            this.cmp = cmp;
+            this.sumType = sumType;
+        }
+
+        /** {@inheritDoc} */
+        @Override public void add(Row row) {
+            boolean first0 = "FIRST".equalsIgnoreCase(get(1, row).toString());
+
+            if (!touched) {
+                touched = true;
+                first = first0;
+                bestRow = row;
+                sum = get(0, row);
+
+                return;
+            }
+
+            assert first == first0;
+
+            int cmp0 = cmp.compare(row, bestRow);
+
+            if ((first && cmp0 < 0) || (!first && cmp0 > 0)) {
+                bestRow = row;
+                sum = get(0, row);
+            }
+            else if (cmp0 == 0)
+                addToSum(get(0, row));
+        }
+
+        /** {@inheritDoc} */
+        @Override public void apply(Accumulator<Row> other) {
+            SumWithKeep<Row> other0 = (SumWithKeep<Row>)other;
+
+            if (!other0.touched)
+                return;
+
+            if (!touched) {
+                touched = true;
+                first = other0.first;
+                bestRow = other0.bestRow;
+                sum = other0.sum;
+
+                return;
+            }
+
+            assert first == other0.first;
+
+            int cmp0 = cmp.compare(other0.bestRow, bestRow);
+
+            if ((first && cmp0 < 0) || (!first && cmp0 > 0)) {
+                bestRow = other0.bestRow;
+                sum = other0.sum;
+            }
+            else if (cmp0 == 0)
+                addToSum(other0.sum);
+        }
+
+        /** */
+        private void addToSum(Object val) {
+            if (val == null)
+                return;
+
+            if (sum == null) {
+                sum = val;
+
+                return;
+            }
+
+            switch (sumType) {
+                case DECIMAL:
+                    sum = ((BigDecimal)sum).add((BigDecimal)val);
+                    break;
+
+                case DOUBLE:
+                    sum = (Double)sum + (Double)val;
+                    break;
+
+                default:
+                    sum = (Long)sum + (Long)val;
+            }
+        }
+
+        /** {@inheritDoc} */
+        @Override public Object end() {
+            return sum;
+        }
+
+        /** {@inheritDoc} */
+        @Override public List<RelDataType> argumentTypes(IgniteTypeFactory typeFactory) {
+            return F.asList(
+                typeFactory.createTypeWithNullability(typeFactory.createSqlType(sumType), true),
+                typeFactory.createTypeWithNullability(typeFactory.createSqlType(VARCHAR), false)
+            );
+        }
+
+        /** {@inheritDoc} */
+        @Override public RelDataType returnType(IgniteTypeFactory typeFactory) {
+            return typeFactory.createTypeWithNullability(typeFactory.createSqlType(sumType), true);
+        }
+    }
+
     /** */
     private static class DoubleSumEmptyIsZero<Row> extends AbstractAccumulator<Row> {
         /** */

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/sql/fun/IgniteOwnSqlOperatorTable.java

@@ -17,6 +17,7 @@
 package org.apache.ignite.internal.processors.query.calcite.sql.fun;
 
 import org.apache.calcite.sql.SqlFunction;
+import org.apache.calcite.sql.SqlAggFunction;
 import org.apache.calcite.sql.SqlFunctionCategory;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.type.InferTypes;
@@ -39,6 +40,9 @@ public class IgniteOwnSqlOperatorTable extends ReflectiveSqlOperatorTable {
      */
     private static IgniteOwnSqlOperatorTable instance;
 
+    /** Sum of values having the first or last ordering key. */
+    public static final SqlAggFunction SUM_WITH_KEEP = new SqlSumWithKeepAggFunction();
+
     /**
      *
      */

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/sql/fun/SqlSumWithKeepAggFunction.java

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.ignite.internal.processors.query.calcite.sql.fun;
+
+import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlFunctionCategory;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.util.Optionality;
+import org.apache.ignite.internal.processors.query.calcite.util.IgniteResource;
+
+import static org.apache.calcite.util.Static.RESOURCE;
+
+/**
+ * Aggregate that sums values having the first or last ordering key.
+ *
+ * <p>The syntax is
+ * {@code SUM_WITH_KEEP(value, 'FIRST'|'LAST') WITHIN GROUP (ORDER BY orderKey [, ...])}.
+ */
+public class SqlSumWithKeepAggFunction extends SqlAggFunction {
+    /** */
+    public SqlSumWithKeepAggFunction() {
+        super(
+            "SUM_WITH_KEEP",
+            null,
+            SqlKind.SUM,
+            ReturnTypes.AGG_SUM,
+            null,
+            OperandTypes.family(SqlTypeFamily.NUMERIC, SqlTypeFamily.CHARACTER),
+            SqlFunctionCategory.NUMERIC,
+            false,
+            false,
+            Optionality.MANDATORY
+        );
+    }
+
+    /** {@inheritDoc} */
+    @Override public boolean checkOperandTypes(SqlCallBinding callBinding, boolean throwOnFailure) {
+        if (!super.checkOperandTypes(callBinding, throwOnFailure))
+            return false;
+
+        if (!callBinding.isOperandLiteral(1, false)) {
+            if (throwOnFailure)
+                throw callBinding.newError(RESOURCE.argumentMustBeLiteral(getName()));
+
+            return false;
+        }
+
+        String mode = callBinding.getStringLiteralOperand(1);
+
+        if (!"FIRST".equalsIgnoreCase(mode) && !"LAST".equalsIgnoreCase(mode)) {
+            if (throwOnFailure)
+                throw callBinding.newError(IgniteResource.INSTANCE.illegalSumWithKeepMode(mode));
+
+            return false;
+        }
+
+        return true;
+    }
+
+    /** {@inheritDoc} */
+    @Override public Optionality getDistinctOptionality() {
+        return Optionality.FORBIDDEN;
+    }
+}

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/util/IgniteResource.java

@@ -39,6 +39,10 @@ public interface IgniteResource {
     @Resources.BaseMessage("Illegal aggregate function. {0} is unsupported at the moment.")
     Resources.ExInst<SqlValidatorException> unsupportedAggregationFunction(String a0);
 
+    /** */
+    @Resources.BaseMessage("Illegal SUM_WITH_KEEP mode ''{0}''. Expected FIRST or LAST.")
+    Resources.ExInst<SqlValidatorException> illegalSumWithKeepMode(String mode);
+
     /** */
     @Resources.BaseMessage("Illegal value of {0}. The value must be positive and less than Integer.MAX_VALUE " +
         "(" + Integer.MAX_VALUE + ")." )

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/exec/rel/BaseAggregateTest.java

@@ -28,6 +28,7 @@
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import com.google.common.collect.ImmutableList;
+import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.type.RelDataType;
@@ -40,6 +41,7 @@
 import org.apache.ignite.internal.processors.query.calcite.exec.exp.agg.AccumulatorWrapper;
 import org.apache.ignite.internal.processors.query.calcite.exec.exp.agg.Accumulators;
 import org.apache.ignite.internal.processors.query.calcite.exec.exp.agg.AggregateType;
+import org.apache.ignite.internal.processors.query.calcite.sql.fun.IgniteOwnSqlOperatorTable;
 import org.apache.ignite.internal.processors.query.calcite.type.IgniteTypeFactory;
 import org.apache.ignite.internal.processors.query.calcite.util.TypeUtils;
 import org.apache.ignite.internal.util.typedef.F;
@@ -242,6 +244,52 @@ public void max() {
         assertFalse(root.hasNext());
     }
 
+    /** */
+    @Test
+    public void sumWithKeep() {
+        ExecutionContext<Object[]> ctx = executionContext(F.first(nodes()), UUID.randomUUID(), 0);
+        IgniteTypeFactory tf = ctx.getTypeFactory();
+        RelDataType rowType = TypeUtils.createRowType(tf, int.class, int.class, String.class, int.class);
+        RelDataType aggRowType = TypeUtils.createRowType(tf, long.class);
+
+        for (String mode : Arrays.asList("FIRST", "LAST")) {
+            ScanNode<Object[]> scan = new ScanNode<>(ctx, rowType, Arrays.asList(
+                row(0, 10, mode, 2),
+                row(0, 20, mode, 1),
+                row(0, 30, mode, 1),
+                row(0, 40, mode, 3)
+            ));
+
+            AggregateCall call = AggregateCall.create(
+                IgniteOwnSqlOperatorTable.SUM_WITH_KEEP,
+                false,
+                false,
+                false,
+                ImmutableIntList.of(1, 2),
+                -1,
+                RelCollations.of(new RelFieldCollation(3)),
+                tf.createJavaType(long.class),
+                null);
+
+            SingleNode<Object[]> aggChain = createAggregateNodesChain(
+                ctx,
+                ImmutableList.of(ImmutableBitSet.of(0)),
+                call,
+                rowType,
+                aggRowType,
+                rowFactory(),
+                scan
+            );
+
+            RootNode<Object[]> root = new RootNode<>(ctx, aggRowType);
+            root.register(aggChain);
+
+            assertTrue(root.hasNext());
+            Assert.assertArrayEquals(row(0, "FIRST".equals(mode) ? 50L : 40L), root.next());
+            assertFalse(root.hasNext());
+        }
+    }
+
     /** */
     @Test
     public void avg() {