feat(snowflake)!:annotation support for APPROX_TOP_K_COMBINE

sqlglot/expressions.py

@@ -5732,6 +5732,11 @@ class ApproxTopKAccumulate(AggFunc):
     arg_types = {"this": True, "expression": False}
 
 
+# https://docs.snowflake.com/en/sql-reference/functions/approx_top_k_combine
+class ApproxTopKCombine(AggFunc):
+    arg_types = {"this": True, "expression": False}
+
+
 class ApproxTopSum(AggFunc):
     arg_types = {"this": True, "expression": True, "count": True}
 

sqlglot/typing/snowflake.py

@@ -244,6 +244,7 @@ def _annotate_within_group(self: TypeAnnotator, expression: exp.WithinGroup) ->
             exp.ParseIp,
             exp.ParseUrl,
             exp.ApproxTopKAccumulate,
+            exp.ApproxTopKCombine,
         }
     },
     **{

tests/dialects/test_snowflake.py

@@ -52,6 +52,7 @@ def test_snowflake(self):
             "SELECT APPROXIMATE_SIMILARITY(minhash_col)",
         )
         self.validate_identity("SELECT APPROX_TOP_K_ACCUMULATE(col, 10)")
+        self.validate_identity("SELECT APPROX_TOP_K_COMBINE(state, 2)")
         self.validate_identity("SELECT EQUAL_NULL(1, 2)")
         self.validate_identity("SELECT EXP(1)")
         self.validate_identity("SELECT FACTORIAL(5)")

tests/fixtures/optimizer/annotate_functions.sql

@@ -3460,6 +3460,14 @@ DOUBLE;
 APPROX_TOP_K_ACCUMULATE(tbl.str_col, 10);
 OBJECT;
 
+# dialect: snowflake
+APPROX_TOP_K_COMBINE(tbl.state_col, 10);
+OBJECT;
+
+# dialect: snowflake
+APPROX_TOP_K_COMBINE(tbl.state_col);
+OBJECT;
+
 # dialect: snowflake
 APPROX_COUNT_DISTINCT(tbl.str_col);
 BIGINT;