feat(snowflake)!: type annotation for APPROX_TOP_K_ACCUMULATE functio…

sqlglot/expressions.py

@@ -5703,6 +5703,10 @@ class ApproxTopK(AggFunc):
     arg_types = {"this": True, "expression": False, "counters": False}
 
 
+class ApproxTopKAccumulate(AggFunc):
+    arg_types = {"this": True, "expression": False}
+
+
 class ApproxTopSum(AggFunc):
     arg_types = {"this": True, "expression": True, "count": True}
 

sqlglot/typing/snowflake.py

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

tests/dialects/test_snowflake.py

@@ -42,6 +42,7 @@ def test_snowflake(self):
         expr.selects[0].assert_is(exp.AggFunc)
         self.assertEqual(expr.sql(dialect="snowflake"), "SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
 
+        self.validate_identity("SELECT APPROX_TOP_K_ACCUMULATE(col, 10)")
         self.validate_identity("SELECT EQUAL_NULL(1, 2)")
         self.validate_identity("SELECT EXP(1)")
         self.validate_identity("SELECT FACTORIAL(5)")

tests/dialects/test_spark.py

@@ -274,6 +274,7 @@ def test_spark(self):
             "REFRESH TABLE t",
         )
 
+        self.validate_identity("SELECT APPROX_TOP_K_ACCUMULATE(col, 10)")
         self.validate_identity("SELECT BITMAP_BIT_POSITION(10)")
         self.validate_identity("SELECT BITMAP_CONSTRUCT_AGG(value)")
         self.validate_identity("ALTER TABLE foo ADD PARTITION(event = 'click')")

tests/fixtures/optimizer/annotate_functions.sql

@@ -3456,6 +3456,10 @@ DOUBLE;
 APPROX_PERCENTILE(tbl.int_col, 0.9) OVER (PARTITION BY 1);
 DOUBLE;
 
+# dialect: snowflake
+APPROX_TOP_K_ACCUMULATE(tbl.str_col, 10);
+OBJECT;
+
 # dialect: snowflake
 APPROX_COUNT_DISTINCT(tbl.str_col);
 BIGINT;