feat(bigquery)!: add typed AI scalar function nodes (#7479)

Signed-off-by: Mridankan Mandal <xerontitan90@gmail.com>
Co-authored-by: Mridankan Mandal <xerontitan90@gmail.com>

Author: RedZapdos123

sqlglot/expressions/functions.py

@@ -321,6 +321,24 @@ class AIClassify(Expression, Func):
     _sql_names = ["AI_CLASSIFY"]
 
 
+class AIEmbed(Expression, Func):
+    arg_types = {"expressions": False}
+    is_var_len_args = True
+    _sql_names = ["EMBED"]
+
+
+class AISimilarity(Expression, Func):
+    arg_types = {"expressions": False}
+    is_var_len_args = True
+    _sql_names = ["SIMILARITY"]
+
+
+class AIGenerate(Expression, Func):
+    arg_types = {"expressions": False}
+    is_var_len_args = True
+    _sql_names = ["GENERATE"]
+
+
 class FeaturesAtTime(Expression, Func):
     arg_types = {"this": True, "time": False, "num_rows": False, "ignore_feature_nulls": False}
 

sqlglot/parsers/bigquery.py

@@ -730,6 +730,16 @@ def _parse_column_ops(self, this: exp.Expr | None) -> exp.Expr | None:
                 self._retreat(func_index)
                 parsed = self._parse_function(any_token=True)
                 if parsed:
+                    if prefix == "AI" and isinstance(parsed, exp.Anonymous):
+                        ai_scalars: dict[str, type[exp.Func]] = {
+                            "EMBED": exp.AIEmbed,
+                            "SIMILARITY": exp.AISimilarity,
+                            "GENERATE": exp.AIGenerate,
+                        }
+                        expr_type = ai_scalars.get(parsed.name.upper())
+                        if expr_type:
+                            parsed = expr_type(expressions=parsed.expressions)
+
                     this = self.expression(exp.Dot(this=this.this, expression=parsed))
 
         return this

tests/dialects/test_bigquery.py

@@ -2449,6 +2449,18 @@ def test_ml_functions(self):
             "SELECT AI.GENERATE_BOOL(MODEL `mydataset.gemini_model`, 'Is sky blue?')"
         )
 
+        ast = self.validate_identity("SELECT AI.EMBED('hello')")
+        assert isinstance(ast.expressions[0], exp.Dot)
+        assert isinstance(ast.expressions[0].expression, exp.AIEmbed)
+
+        ast = self.validate_identity("SELECT AI.SIMILARITY('a', 'b')")
+        assert isinstance(ast.expressions[0], exp.Dot)
+        assert isinstance(ast.expressions[0].expression, exp.AISimilarity)
+
+        ast = self.validate_identity("SELECT AI.GENERATE('Write a haiku')")
+        assert isinstance(ast.expressions[0], exp.Dot)
+        assert isinstance(ast.expressions[0].expression, exp.AIGenerate)
+
     def test_merge(self):
         self.validate_all(
             """