Allow optional second argument (max map size) for FREQUENTSTRINGSSKET…

…CH / FREQUENTLONGSSKETCH in the multi-stage engine (#14939)
apache · Jan 29, 2025 · 2f4a2d5 · 2f4a2d5
1 parent 3525116
commit 2f4a2d5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java
@@ -257,7 +257,11 @@ protected Object[][] provideQueries() {
         new Object[]{"SELECT ts_timestamp - CAST(123456789 AS TIMESTAMP) FROM a"},
         new Object[]{"SELECT SUB(ts_timestamp, CAST(123456789 AS TIMESTAMP)) FROM a"},
         new Object[]{"SELECT ts_timestamp + CAST(123456789 AS TIMESTAMP) FROM a"},
-        new Object[]{"SELECT ADD(ts_timestamp, CAST(123456789 AS TIMESTAMP)) FROM a"}
+        new Object[]{"SELECT ADD(ts_timestamp, CAST(123456789 AS TIMESTAMP)) FROM a"},
+        new Object[]{"SELECT FREQUENT_STRINGS_SKETCH(col1, 512) FROM a"},
+        new Object[]{"SELECT FREQUENT_STRINGS_SKETCH(col1) FROM a"},
+        new Object[]{"SELECT FREQUENT_LONGS_SKETCH(col3, 1024) FROM a"},
+        new Object[]{"SELECT FREQUENT_LONGS_SKETCH(col3) FROM a"},
     };
   }
 

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
@@ -142,8 +142,10 @@ public enum AggregationFunctionType {
   FOURTHMOMENT("fourthMoment", ReturnTypes.DOUBLE, OperandTypes.ANY, SqlTypeName.OTHER),
 
   // Datasketches Frequent Items support
-  FREQUENTSTRINGSSKETCH("frequentStringsSketch", ReturnTypes.VARCHAR, OperandTypes.ANY, SqlTypeName.OTHER),
-  FREQUENTLONGSSKETCH("frequentLongsSketch", ReturnTypes.VARCHAR, OperandTypes.ANY, SqlTypeName.OTHER),
+  FREQUENTSTRINGSSKETCH("frequentStringsSketch", ReturnTypes.VARCHAR,
+      OperandTypes.family(List.of(SqlTypeFamily.ANY, SqlTypeFamily.INTEGER), i -> i == 1), SqlTypeName.OTHER),
+  FREQUENTLONGSSKETCH("frequentLongsSketch", ReturnTypes.VARCHAR,
+      OperandTypes.family(List.of(SqlTypeFamily.ANY, SqlTypeFamily.INTEGER), i -> i == 1), SqlTypeName.OTHER),
 
   // Geo aggregation functions
   STUNION("STUnion", ReturnTypes.VARBINARY, OperandTypes.BINARY, SqlTypeName.OTHER),