feat(optimizer)!: annotate type for Snowflake REGEXP_COUNT function (#5963)

Author: fivetran-BradfordPaskewitz

sqlglot/dialects/snowflake.py

@@ -590,6 +590,9 @@ class Snowflake(Dialect):
             exp.ParseUrl,
             exp.ParseIp,
         },
+        exp.DataType.Type.DECIMAL: {
+            exp.RegexpCount,
+        },
     }
 
     ANNOTATORS = {
@@ -611,6 +614,9 @@ class Snowflake(Dialect):
         },
         exp.ConcatWs: lambda self, e: self._annotate_by_args(e, "expressions"),
         exp.Reverse: _annotate_reverse,
+        exp.RegexpCount: lambda self, e: self._annotate_with_type(
+            e, exp.DataType.build("NUMBER", dialect="snowflake")
+        ),
     }
 
     TIME_MAPPING = {

sqlglot/expressions.py

@@ -7293,6 +7293,15 @@ class RegexpSplit(Func):
     arg_types = {"this": True, "expression": True, "limit": False}
 
 
+class RegexpCount(Func):
+    arg_types = {
+        "this": True,
+        "expression": True,
+        "position": False,
+        "parameters": False,
+    }
+
+
 class Repeat(Func):
     arg_types = {"this": True, "times": True}
 

tests/dialects/test_snowflake.py

@@ -2421,6 +2421,10 @@ def test_regexp_substr(self, logger):
             "REGEXP_EXTRACT_ALL(subject, pattern)",
         )
 
+        self.validate_identity("SELECT REGEXP_COUNT('hello world', 'l')")
+        self.validate_identity("SELECT REGEXP_COUNT('hello world', 'l', 1)")
+        self.validate_identity("SELECT REGEXP_COUNT('hello world', 'l', 1, 'i')")
+
     @mock.patch("sqlglot.generator.logger")
     def test_regexp_replace(self, logger):
         self.validate_all(

tests/fixtures/optimizer/annotate_functions.sql

@@ -1883,6 +1883,18 @@ BOOLEAN;
 REGEXP_LIKE('foo', NULL, 'baz');
 BOOLEAN;
 
+# dialect: snowflake
+REGEXP_COUNT('hello world', 'l');
+DECIMAL(38, 0);
+
+# dialect: snowflake
+REGEXP_COUNT('hello world', 'l', 1);
+DECIMAL(38, 0);
+
+# dialect: snowflake
+REGEXP_COUNT('hello world', 'l', 1, 'i');
+DECIMAL(38, 0);
+
 # dialect: snowflake
 REGEXP_REPLACE('hello world', 'world', 'universe');
 VARCHAR;