Add LASTWITHTIME aggregate function support #7315

pinot-common/src/test/java/org/apache/pinot/common/function/AggregationFunctionTypeTest.java

@@ -33,6 +33,8 @@ public void testGetAggregationFunctionType() {
     Assert.assertEquals(AggregationFunctionType.getAggregationFunctionType("SuM"), AggregationFunctionType.SUM);
     Assert.assertEquals(AggregationFunctionType.getAggregationFunctionType("AvG"), AggregationFunctionType.AVG);
     Assert.assertEquals(AggregationFunctionType.getAggregationFunctionType("MoDe"), AggregationFunctionType.MODE);
+    Assert.assertEquals(AggregationFunctionType.getAggregationFunctionType("LaStWiThTiMe"),
+            AggregationFunctionType.LASTWITHTIME);
     Assert.assertEquals(AggregationFunctionType.getAggregationFunctionType("MiNmAxRaNgE"),
         AggregationFunctionType.MINMAXRANGE);
     Assert.assertEquals(AggregationFunctionType.getAggregationFunctionType("DiStInCtCoUnT"),

pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java

@@ -63,6 +63,7 @@
 import org.apache.pinot.core.query.utils.idset.IdSet;
 import org.apache.pinot.core.query.utils.idset.IdSets;
 import org.apache.pinot.segment.local.customobject.AvgPair;
+import org.apache.pinot.segment.local.customobject.LastWithTimePair;
 import org.apache.pinot.segment.local.customobject.MinMaxRangePair;
 import org.apache.pinot.segment.local.customobject.QuantileDigest;
 import org.apache.pinot.segment.local.utils.GeometrySerializer;
@@ -109,7 +110,8 @@ public enum ObjectType {
     Int2LongMap(23),
     Long2LongMap(24),
     Float2LongMap(25),
-    Double2LongMap(26);
+    Double2LongMap(26),
+    LastWithTimePair(27);
     private final int _value;
 
     ObjectType(int value) {
@@ -178,6 +180,8 @@ public static ObjectType getObjectType(Object value) {
         return ObjectType.IdSet;
       } else if (value instanceof List) {
         return ObjectType.List;
+      } else if (value instanceof LastWithTimePair) {
+        return ObjectType.LastWithTimePair;
       } else {
         throw new IllegalArgumentException("Unsupported type of value: " + value.getClass().getSimpleName());
       }
@@ -330,6 +334,24 @@ public MinMaxRangePair deserialize(ByteBuffer byteBuffer) {
     }
   };
 
+  public static final ObjectSerDe<LastWithTimePair> LAST_WITH_TIME_PAIR_SER_DE = new ObjectSerDe<LastWithTimePair>() {
+
+    @Override
+    public byte[] serialize(LastWithTimePair lastWithTimePair) {
+      return lastWithTimePair.toBytes();
+    }
+
+    @Override
+    public LastWithTimePair deserialize(byte[] bytes) {
+      return LastWithTimePair.fromBytes(bytes);
+    }
+
+    @Override
+    public LastWithTimePair deserialize(ByteBuffer byteBuffer) {
+      return LastWithTimePair.fromByteBuffer(byteBuffer);
+    }
+  };
+
   public static final ObjectSerDe<HyperLogLog> HYPER_LOG_LOG_SER_DE = new ObjectSerDe<HyperLogLog>() {
 
     @Override
@@ -1047,7 +1069,8 @@ public Double2LongOpenHashMap deserialize(ByteBuffer byteBuffer) {
       INT_2_LONG_MAP_SER_DE,
       LONG_2_LONG_MAP_SER_DE,
       FLOAT_2_LONG_MAP_SER_DE,
-      DOUBLE_2_LONG_MAP_SER_DE
+      DOUBLE_2_LONG_MAP_SER_DE,
+      LAST_WITH_TIME_PAIR_SER_DE
   };
   //@formatter:on
 

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java

@@ -156,6 +156,13 @@ public static AggregationFunction getAggregationFunction(FunctionContext functio
             return new AvgAggregationFunction(firstArgument);
           case MODE:
             return new ModeAggregationFunction(arguments);
+          case LASTWITHTIME:
+            if (arguments.size() > 1) {
+              ExpressionContext secondArgument = arguments.get(1);
+              return new LastWithTimeAggregationFunction(firstArgument, secondArgument);
+            } else {
+              throw new IllegalArgumentException();
+            }
           case MINMAXRANGE:
             return new MinMaxRangeAggregationFunction(firstArgument);
           case DISTINCTCOUNT:

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountAggregationFunction.java

@@ -42,7 +42,6 @@
 
 @SuppressWarnings({"rawtypes", "unchecked"})
 public class DistinctCountAggregationFunction extends BaseSingleInputAggregationFunction<Set, Integer> {
-
   public DistinctCountAggregationFunction(ExpressionContext expression) {
     super(expression);
   }

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/LastWithTimeAggregationFunction.java

@@ -0,0 +1,222 @@
+/**
+ * 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.pinot.core.query.aggregation.function;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.core.common.BlockValSet;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
+import org.apache.pinot.core.query.aggregation.ObjectAggregationResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
+import org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder;
+import org.apache.pinot.segment.local.customobject.LastWithTimePair;
+import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+public class LastWithTimeAggregationFunction extends BaseSingleInputAggregationFunction<LastWithTimePair, Double> {
+  private final ExpressionContext _timeCol;
+
+  public LastWithTimeAggregationFunction(ExpressionContext dataCol, ExpressionContext timeCol) {
+    super(dataCol);
+    _timeCol = timeCol;
+  }
+
+  @Override
+  public List<ExpressionContext> getInputExpressions() {
+    return Arrays.asList(_expression, _timeCol);
+  }
+
+  @Override
+  public AggregationFunctionType getType() {
+    return AggregationFunctionType.LASTWITHTIME;
+  }
+
+  @Override
+  public AggregationResultHolder createAggregationResultHolder() {
+    return new ObjectAggregationResultHolder();
+  }
+
+  @Override
+  public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int maxCapacity) {
+    return new ObjectGroupByResultHolder(initialCapacity, maxCapacity);
+  }
+
+  @Override
+  public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
+      Map<ExpressionContext, BlockValSet> blockValSetMap) {
+    double lastData = Double.POSITIVE_INFINITY;
+    long lastTime = Long.MIN_VALUE;
+
+    BlockValSet blockValSet = blockValSetMap.get(_expression);
+    BlockValSet blockTimeSet = blockValSetMap.get(_timeCol);
+    if (blockValSet.getValueType() != DataType.BYTES) {
+      double[] doubleValues = blockValSet.getDoubleValuesSV();
+      long[] timeValues = blockTimeSet.getLongValuesSV();
+      for (int i = 0; i < length; i++) {
+        double data = doubleValues[i];
+        long time = timeValues[i];
+        if (time >= lastTime) {
+          lastTime = time;
+          lastData = data;
+        }
+      }
+    } else {
+      // Serialized LastPair
+      byte[][] bytesValues = blockValSet.getBytesValuesSV();
+      for (int i = 0; i < length; i++) {
+        LastWithTimePair lastWithTimePair = ObjectSerDeUtils.LAST_WITH_TIME_PAIR_SER_DE.deserialize(bytesValues[i]);
+        double data = lastWithTimePair.getData();
+        long time = lastWithTimePair.getTime();
+        if (time >= lastTime) {
+          lastTime = time;
+          lastData = data;
+        }
+      }
+    }
+    setAggregationResult(aggregationResultHolder, lastData, lastTime);
+  }
+
+  protected void setAggregationResult(AggregationResultHolder aggregationResultHolder, double data, long time) {
+    LastWithTimePair lastWithTimePair = aggregationResultHolder.getResult();
+    if (lastWithTimePair == null) {
+      aggregationResultHolder.setValue(new LastWithTimePair(data, time));
+    } else {
+      lastWithTimePair.apply(data, time);
+    }
+  }
+
+  @Override
+  public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
+      Map<ExpressionContext, BlockValSet> blockValSetMap) {
+    BlockValSet blockValSet = blockValSetMap.get(_expression);
+    BlockValSet timeValSet = blockValSetMap.get(_timeCol);
+    if (blockValSet.getValueType() != DataType.BYTES) {
+      double[] doubleValues = blockValSet.getDoubleValuesSV();
+      long[] timeValues = timeValSet.getLongValuesSV();
+      for (int i = 0; i < length; i++) {
+        double data = doubleValues[i];
+        long time = timeValues[i];
+        setGroupByResult(groupKeyArray[i], groupByResultHolder, data, time);
+      }
+    } else {
+      // Serialized LastPair
+      byte[][] bytesValues = blockValSet.getBytesValuesSV();
+      for (int i = 0; i < length; i++) {
+        LastWithTimePair lastWithTimePair = ObjectSerDeUtils.LAST_WITH_TIME_PAIR_SER_DE.deserialize(bytesValues[i]);
+        setGroupByResult(groupKeyArray[i], groupByResultHolder, lastWithTimePair.getData(), lastWithTimePair.getTime());
+      }
+    }
+  }
+
+  @Override
+  public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
+      Map<ExpressionContext, BlockValSet> blockValSetMap) {
+    BlockValSet blockValSet = blockValSetMap.get(_expression);
+    BlockValSet timeValSet = blockValSetMap.get(_timeCol);
+    if (blockValSet.getValueType() != DataType.BYTES) {
+      double[] doubleValues = blockValSet.getDoubleValuesSV();
+      long[] timeValues = timeValSet.getLongValuesSV();
+      for (int i = 0; i < length; i++) {
+        double value = doubleValues[i];
+        long time = timeValues[i];
+        for (int groupKey : groupKeysArray[i]) {
+          setGroupByResult(groupKey, groupByResultHolder, value, time);
+        }
+      }
+    } else {
+      // Serialized LastWithTimePair
+      byte[][] bytesValues = blockValSet.getBytesValuesSV();
+      for (int i = 0; i < length; i++) {
+        LastWithTimePair lastWithTimePair = ObjectSerDeUtils.LAST_WITH_TIME_PAIR_SER_DE.deserialize(bytesValues[i]);
+        double data = lastWithTimePair.getData();
+        long time = lastWithTimePair.getTime();
+        for (int groupKey : groupKeysArray[i]) {
+          setGroupByResult(groupKey, groupByResultHolder, data, time);
+        }
+      }
+    }
+  }
+
+  protected void setGroupByResult(int groupKey, GroupByResultHolder groupByResultHolder, double data, long time) {
+    LastWithTimePair lastWithTimePair = groupByResultHolder.getResult(groupKey);
+    if (lastWithTimePair == null) {
+      groupByResultHolder.setValueForKey(groupKey, new LastWithTimePair(data, time));
+    } else {
+      lastWithTimePair.apply(data, time);
+    }
+  }
+
+  @Override
+  public LastWithTimePair extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
+    LastWithTimePair lastWithTimePair = aggregationResultHolder.getResult();
+    if (lastWithTimePair == null) {
+      return new LastWithTimePair(Double.POSITIVE_INFINITY, Long.MIN_VALUE);
+    } else {
+      return lastWithTimePair;
+    }
+  }
+
+  @Override
+  public LastWithTimePair extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
+    LastWithTimePair lastWithTimePair = groupByResultHolder.getResult(groupKey);
+    if (lastWithTimePair == null) {
+      return new LastWithTimePair(Double.POSITIVE_INFINITY, Long.MIN_VALUE);
+    } else {
+      return lastWithTimePair;
+    }
+  }
+
+  @Override
+  public LastWithTimePair merge(LastWithTimePair intermediateResult1, LastWithTimePair intermediateResult2) {
+    intermediateResult1.apply(intermediateResult2);
+    return intermediateResult1;
+  }
+
+  @Override
+  public boolean isIntermediateResultComparable() {
+    return true;
+  }
+
+  @Override
+  public ColumnDataType getIntermediateResultColumnType() {
+    return ColumnDataType.OBJECT;
+  }
+
+  @Override
+  public ColumnDataType getFinalResultColumnType() {
+    return ColumnDataType.DOUBLE;
+  }
+
+  @Override
+  public Double extractFinalResult(LastWithTimePair intermediateResult) {
+    return intermediateResult.getData();
+  }
+
+  @Override
+  public String getResultColumnName() {
+    return getType().getName().toLowerCase() + "(" + _expression + "," + _timeCol + ")";
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverter.java

@@ -159,7 +159,8 @@ private static QueryContext convertPQL(BrokerRequest brokerRequest) {
         List<String> stringExpressions = aggregationInfo.getExpressions();
         int numArguments = stringExpressions.size();
         List<ExpressionContext> arguments = new ArrayList<>(numArguments);
-        if (functionName.equalsIgnoreCase(AggregationFunctionType.DISTINCT.getName())) {
+        if (functionName.equalsIgnoreCase(AggregationFunctionType.DISTINCT.getName())
+          || functionName.equalsIgnoreCase(AggregationFunctionType.LASTWITHTIME.getName())) {
           // For DISTINCT query, all arguments are expressions
           for (String expression : stringExpressions) {
             arguments.add(RequestContextUtils.getExpressionFromPQL(expression));

pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsTest.java

@@ -34,6 +34,7 @@
 import org.apache.pinot.core.query.aggregation.function.PercentileEstAggregationFunction;
 import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction;
 import org.apache.pinot.segment.local.customobject.AvgPair;
+import org.apache.pinot.segment.local.customobject.LastWithTimePair;
 import org.apache.pinot.segment.local.customobject.MinMaxRangePair;
 import org.apache.pinot.segment.local.customobject.QuantileDigest;
 import org.testng.annotations.Test;
@@ -126,6 +127,19 @@ public void testMinMaxRangePair() {
     }
   }
 
+  @Test
+  public void testLastWithTimePair() {
+    for (int i = 0; i < NUM_ITERATIONS; i++) {
+      LastWithTimePair expected = new LastWithTimePair(RANDOM.nextDouble(), RANDOM.nextLong());
+
+      byte[] bytes = ObjectSerDeUtils.serialize(expected);
+      LastWithTimePair actual = ObjectSerDeUtils.deserialize(bytes, ObjectSerDeUtils.ObjectType.LastWithTimePair);
+
+      assertEquals(actual.getData(), expected.getData(), ERROR_MESSAGE);
+      assertEquals(actual.getTime(), expected.getTime(), ERROR_MESSAGE);
+    }
+  }
+
   @Test
   public void testHyperLogLog() {
     for (int i = 0; i < NUM_ITERATIONS; i++) {

pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactoryTest.java

@@ -32,6 +32,7 @@
 @SuppressWarnings("rawtypes")
 public class AggregationFunctionFactoryTest {
   private static final String ARGUMENT = "(column)";
+  private static final String ARGUMENTS = "(column,timeColumn)";
   private static final QueryContext DUMMY_QUERY_CONTEXT =
       QueryContextConverterUtils.getQueryContextFromSQL("SELECT * FROM testTable");
 
@@ -87,6 +88,13 @@ public void testGetAggregationFunction() {
     assertEquals(aggregationFunction.getColumnName(), "mode_column");
     assertEquals(aggregationFunction.getResultColumnName(), function.toString());
 
+    function = getFunctionWithTwoArgs("LaStWiThTiMe");
+    aggregationFunction = AggregationFunctionFactory.getAggregationFunction(function, DUMMY_QUERY_CONTEXT);
+    assertTrue(aggregationFunction instanceof LastWithTimeAggregationFunction);
+    assertEquals(aggregationFunction.getType(), AggregationFunctionType.LASTWITHTIME);
+    assertEquals(aggregationFunction.getColumnName(), "lastWithTime_column");
+    assertEquals(aggregationFunction.getResultColumnName(), function.toString());
+
     function = getFunction("MiNmAxRaNgE");
     aggregationFunction = AggregationFunctionFactory.getAggregationFunction(function, DUMMY_QUERY_CONTEXT);
     assertTrue(aggregationFunction instanceof MinMaxRangeAggregationFunction);
@@ -365,6 +373,10 @@ private FunctionContext getFunction(String functionName) {
     return getFunction(functionName, ARGUMENT);
   }
 
+  private FunctionContext getFunctionWithTwoArgs(String functionName) {
+    return getFunction(functionName, ARGUMENTS);
+  }
+
   private FunctionContext getFunction(String functionName, String args) {
     return RequestContextUtils.getExpressionFromSQL(functionName + args).getFunction();
   }