feat(typing): add databricks type inference module for REGEXP_COUNT, REGEXP_EXTRACT_ALL

sqlglot/dialects/databricks.py

@@ -9,11 +9,13 @@
 from sqlglot.parsers.databricks import DatabricksParser
 from sqlglot.tokens import TokenType
 from sqlglot.optimizer.annotate_types import TypeAnnotator
+from sqlglot.typing.databricks import EXPRESSION_METADATA
 
 
 class Databricks(Spark):
     SAFE_DIVISION = False
     COPY_PARAMS_ARE_CSV = False
+    EXPRESSION_METADATA = EXPRESSION_METADATA.copy()
 
     COERCES_TO = defaultdict(set, deepcopy(TypeAnnotator.COERCES_TO))
     for text_type in exp.DataType.TEXT_TYPES:

sqlglot/typing/databricks.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from sqlglot import exp
+from sqlglot.typing.spark import EXPRESSION_METADATA
+
+EXPRESSION_METADATA = {
+    **EXPRESSION_METADATA,
+    **{
+        exp_type: {"returns": exp.DType.INT}
+        for exp_type in {
+            exp.RegexpCount,
+        }
+    },
+    **{
+        exp_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this", array=True)}
+        for exp_type in {
+            exp.RegexpExtractAll,
+        }
+    },
+}

tests/fixtures/optimizer/annotate_functions.sql

@@ -970,7 +970,27 @@ BIGINT;
 
 # dialect: spark2, spark, databricks
 tbl.double_col DIV tbl.double_col;
-BIGINT; 
+BIGINT;
+
+# dialect: databricks
+tbl.str_col REGEXP 'pattern';
+BOOLEAN;
+
+# dialect: databricks
+tbl.str_col REGEXP tbl.str_col;
+BOOLEAN;
+
+# dialect: databricks
+REGEXP_COUNT(tbl.str_col, 'l');
+INT;
+
+# dialect: databricks
+REGEXP_COUNT(tbl.str_col, tbl.str_col);
+INT;
+
+# dialect: databricks
+REGEXP_EXTRACT_ALL(tbl.str_col, 'pattern');
+ARRAY<VARCHAR>;
 
 # dialect: hive
 tbl.bigint DIV tbl.bigint;