[Enhancement] Add Trino HLL Function Compatibility Mapping and last_d…

…ay_of_month Support(StarRocks#40894) Signed-off-by: happut <[email protected]>
happut · Jun 17, 2024 · fdd4821 · fdd4821
1 parent b0e6bb3
commit fdd4821
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git a/...e/src/main/java/com/starrocks/connector/parser/trino/Trino2SRFunctionCallTransformer.java b/...e/src/main/java/com/starrocks/connector/parser/trino/Trino2SRFunctionCallTransformer.java
@@ -81,6 +81,7 @@ private static void registerAllFunctionTransformer() {
         registerUnicodeFunctionTransformer();
         registerMapFunctionTransformer();
         registerBinaryFunctionTransformer();
+        registerHLLFunctionTransformer();
         // todo: support more function transform
     }
 
@@ -195,6 +196,10 @@ private static void registerDateFunctionTransformer() {
         // to_timestamp -> to_tera_timestamp
         registerFunctionTransformer("to_timestamp", 2, "to_tera_timestamp",
                 ImmutableList.of(Expr.class, Expr.class));
+
+        // last_day_of_month(x)  -> last_day(x,'month')
+        registerFunctionTransformer("last_day_of_month", 1, new FunctionCallExpr("last_day",
+                ImmutableList.of(new PlaceholderExpr(1, Expr.class), new StringLiteral("month"))));
     }
 
     private static void registerStringFunctionTransformer() {
@@ -311,12 +316,31 @@ private static void registerBinaryFunctionTransformer() {
         registerFunctionTransformer("from_hex", 1, "hex_decode_binary", ImmutableList.of(Expr.class));
     }
 
+    private static void registerHLLFunctionTransformer() {
+        // cardinality -> hll_cardinality
+        registerFunctionTransformer("cardinality", 1, "hll_cardinality", ImmutableList.of(Expr.class));
+
+        // approx_set -> HLL_HASH
+        registerFunctionTransformer("approx_set", 1, "hll_hash", ImmutableList.of(Expr.class));
+
+        // empty_approx_set -> HLL_EMPTY
+        registerFunctionTransformer("empty_approx_set", "hll_empty");
+
+        // merge -> HLL_RAW_AGG
+        registerFunctionTransformer("merge", 1, "hll_raw_agg", ImmutableList.of(Expr.class));
+    }
+
     private static void registerFunctionTransformer(String trinoFnName, int trinoFnArgNums, String starRocksFnName,
                                                     List<Class<? extends Expr>> starRocksArgumentsClass) {
         FunctionCallExpr starRocksFunctionCall = buildStarRocksFunctionCall(starRocksFnName, starRocksArgumentsClass);
         registerFunctionTransformer(trinoFnName, trinoFnArgNums, starRocksFunctionCall);
     }
 
+    private static void registerFunctionTransformer(String trinoFnName, String starRocksFnName) {
+        FunctionCallExpr starRocksFunctionCall = buildStarRocksFunctionCall(starRocksFnName, Lists.newArrayList());
+        registerFunctionTransformer(trinoFnName, 0, starRocksFunctionCall);
+    }
+
     private static void registerFunctionTransformerWithVarArgs(String trinoFnName, String starRocksFnName,
                                                                List<Class<? extends Expr>> starRocksArgumentsClass) {
         Preconditions.checkState(starRocksArgumentsClass.size() == 1);
@@ -326,7 +350,12 @@ private static void registerFunctionTransformerWithVarArgs(String trinoFnName, S
 
     private static void registerFunctionTransformer(String trinoFnName, int trinoFnArgNums,
                                                     FunctionCallExpr starRocksFunctionCall) {
-        FunctionCallTransformer transformer = new FunctionCallTransformer(starRocksFunctionCall, trinoFnArgNums);
+        FunctionCallTransformer transformer;
+        if (trinoFnArgNums == 0) {
+            transformer = new FunctionCallTransformer(starRocksFunctionCall, false);
+        } else {
+            transformer = new FunctionCallTransformer(starRocksFunctionCall, trinoFnArgNums);
+        }
 
         List<FunctionCallTransformer> transformerList = TRANSFORMER_MAP.computeIfAbsent(trinoFnName,
                 k -> Lists.newArrayList());

diff --git a/...e-core/src/test/java/com/starrocks/connector/parser/trino/TrinoFunctionTransformTest.java b/...e-core/src/test/java/com/starrocks/connector/parser/trino/TrinoFunctionTransformTest.java
@@ -441,4 +441,20 @@ public void testUtilityFunction() throws Exception {
         sql = "select current_schema";
         assertPlanContains(sql, "<slot 2> : 'test'");
     }
+
+
+    @Test
+    public void testHllFunction() throws Exception {
+        String sql = "select empty_approx_set()";
+        assertPlanContains(sql, "<slot 2> : HLL_EMPTY()");
+
+        sql = "select cardinality(empty_approx_set())";
+        assertPlanContains(sql, "<slot 2> : hll_cardinality(HLL_EMPTY())");
+
+        sql = "select approx_set(\"tc\") from tall";
+        assertPlanContains(sql, "<slot 12> : hll_hash(CAST(3: tc AS VARCHAR))");
+
+        sql = "select merge(approx_set(\"tc\")) from tall";
+        assertPlanContains(sql, "hll_raw_agg(hll_hash(CAST(3: tc AS VARCHAR)))");
+    }
 }