Loose type check for FirstWithTime and LastWithTime (#11683)

* Loose type check for FirstWithTime and LastWithTime * Make type check more strict
apache · Sep 27, 2023 · e686b6a · e686b6a
1 parent 8ce3062
commit e686b6a
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 6 deletions.
diff --git a/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TimestampTest.java b/...egration-tests/src/test/java/org/apache/pinot/integration/tests/custom/TimestampTest.java
@@ -370,6 +370,30 @@ public void testToDateTimeQueries(boolean useMultiStageQueryEngine)
     assertEquals(jsonNode.get("resultTable").get("rows").get(0).get(1).textValue(), "2019-01-01 12:00:00");
   }
 
+  @Test(dataProvider = "useBothQueryEngines")
+  public void testLastWithTimeQueries(boolean useMultiStageQueryEngine)
+      throws Exception {
+    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+    String query = String.format(
+        "SELECT LASTWITHTIME(longBase, longBase, 'long'), LASTWITHTIME(longBase, tsBase, 'long') FROM %s\n",
+        getTableName());
+    JsonNode jsonNode = postQuery(query);
+    assertEquals(jsonNode.get("resultTable").get("rows").get(0).get(0).longValue(), 1632614400000L);
+    assertEquals(jsonNode.get("resultTable").get("rows").get(0).get(1).longValue(), 1632614400000L);
+  }
+
+  @Test(dataProvider = "useBothQueryEngines")
+  public void testFirstWithTimeQueries(boolean useMultiStageQueryEngine)
+      throws Exception {
+    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+    String query = String.format(
+        "SELECT FIRSTWITHTIME(longBase, longBase, 'long'), FIRSTWITHTIME(longBase, tsBase, 'long') FROM %s\n",
+        getTableName());
+    JsonNode jsonNode = postQuery(query);
+    assertEquals(jsonNode.get("resultTable").get("rows").get(0).get(0).longValue(), 1546300800000L);
+    assertEquals(jsonNode.get("resultTable").get("rows").get(0).get(1).longValue(), 1546300800000L);
+  }
+
   @Override
   public String getTableName() {
     return DEFAULT_TABLE_NAME;

diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/AggregationFunctionType.java
@@ -66,12 +66,17 @@ public enum AggregationFunctionType {
   // NO NEEDED in v2, AVG is compiled as SUM/COUNT
   AVG("avg"),
   MODE("mode"),
-
   FIRSTWITHTIME("firstWithTime", null, SqlKind.OTHER_FUNCTION, SqlFunctionCategory.USER_DEFINED_FUNCTION,
-      OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER)),
+      OperandTypes.or(
+          OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER)),
+          OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.NUMERIC, SqlTypeFamily.CHARACTER))
+      ),
       ReturnTypes.ARG0, ReturnTypes.explicit(SqlTypeName.OTHER)),
   LASTWITHTIME("lastWithTime", null, SqlKind.OTHER_FUNCTION, SqlFunctionCategory.USER_DEFINED_FUNCTION,
-      OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER)),
+      OperandTypes.or(
+          OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.TIMESTAMP, SqlTypeFamily.CHARACTER)),
+          OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.NUMERIC, SqlTypeFamily.CHARACTER))
+      ),
       ReturnTypes.ARG0, ReturnTypes.explicit(SqlTypeName.OTHER)),
   MINMAXRANGE("minMaxRange", null, SqlKind.OTHER_FUNCTION, SqlFunctionCategory.NUMERIC,
       OperandTypes.NUMERIC, ReturnTypes.ARG0, ReturnTypes.explicit(SqlTypeName.OTHER)),
@@ -143,13 +148,13 @@ public enum AggregationFunctionType {
       ReturnTypes.explicit(SqlTypeName.OTHER)),
   // hyper log log plus plus functions
   DISTINCTCOUNTHLLPLUS("distinctCountHLLPlus", ImmutableList.of("DISTINCT_COUNT_HLL_PLUS"), SqlKind.OTHER_FUNCTION,
-  SqlFunctionCategory.USER_DEFINED_FUNCTION,
+      SqlFunctionCategory.USER_DEFINED_FUNCTION,
       OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.NUMERIC), ordinal -> ordinal > 0),
-  ReturnTypes.BIGINT, ReturnTypes.explicit(SqlTypeName.OTHER)),
+      ReturnTypes.BIGINT, ReturnTypes.explicit(SqlTypeName.OTHER)),
   DISTINCTCOUNTRAWHLLPLUS("distinctCountRawHLLPlus", ImmutableList.of("DISTINCT_COUNT_RAW_HLL_PLUS"),
       SqlKind.OTHER_FUNCTION, SqlFunctionCategory.USER_DEFINED_FUNCTION,
       OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.INTEGER), ordinal -> ordinal > 0),
-  ReturnTypes.VARCHAR_2000, ReturnTypes.explicit(SqlTypeName.OTHER)),
+      ReturnTypes.VARCHAR_2000, ReturnTypes.explicit(SqlTypeName.OTHER)),
 
   // DEPRECATED in v2
   @Deprecated