Repository: tobymao/sqlglot

Issue:
Missing Support : Google BigQuery ML functions 

# Description
 I'm encountering an issue with the BigQuery dialect in SQLGlot where it fails to parse queries containing BigQuery ML functions like ML.FEATURES_AT_TIME. The expected behavior is to convert the BigQuery query to an AST, but SQLGlot does not currently support these functions, resulting in a parse failure.

## Details
BigQuery ML provides specialized functions to interact with machine learning models and feature tables. One such function, ML.FEATURES_AT_TIME, is used to retrieve the values of features for entities at a given point in time. However, the BigQuery dialect in SQLGlot does not currently recognize or support this function's syntax, leading to parsing errors.

The specific functions and syntax that are not supported include:

``` ML.FEATURES_AT_TIME Function ```
Syntax:

``` 
ML.FEATURES_AT_TIME(
  TABLE feature_table_name,
  time => timestamp_expression,
  [num_rows => int_expression],
  [ignore_feature_nulls => boolean_expression]
)
```
Parameters:
#### feature_table_name: The name of the feature table (e.g., a table containing entity IDs, feature values, and timestamps).
#### time: A specific point in time for which the feature values are retrieved.
#### num_rows (Optional): The maximum number of rows to return for each entity.
#### ignore_feature_nulls (Optional): A flag to determine if rows with null feature values should be ignored.
Example Query:
```
WITH
  feature_table AS (
    SELECT * FROM UNNEST(
      ARRAY<STRUCT<entity_id STRING, f_1 FLOAT64, f_2 FLOAT64, feature_timestamp TIMESTAMP>>[
        ('id1', 1.0, 1.0, TIMESTAMP '2022-06-10 12:00:00+00'),
        ('id2', 12.0, 24.0, TIMESTAMP '2022-06-11 12:00:00+00'),
        ('id1', 11.0, NULL, TIMESTAMP '2022-06-11 12:00:00+00'),
        ('id1', 6.0, 12.0, TIMESTAMP '2022-06-11 10:00:00+00'),
        ('id2', 2.0, 4.0, TIMESTAMP '2022-06-10 12:00:00+00'),
        ('id2', 7.0, NULL, TIMESTAMP '2022-06-11 10:00:00+00')])
  )
SELECT *
FROM
  ML.FEATURES_AT_TIME(
    TABLE feature_table,
    time => '2022-06-12 10:00:00+00',
    num_rows => 1,
    ignore_feature_nulls => TRUE);
```
## Use Case:
#### This function allows feature engineering for machine learning models by fetching specific feature values at a given timestamp for various entities.
### Expected Behavior
#### SQLGlot should parse BigQuery queries containing ML.FEATURES_AT_TIME without errors.
#### The resulting AST should represent the query structure, enabling transformations, validation, or migrations.
### Actual Behavior
#### SQLGlot fails to recognize ML.FEATURES_AT_TIME, resulting in a parsing error.
#### Steps to Reproduce
##### Write a BigQuery query with the ML.FEATURES_AT_TIME function as shown in the example above.
##### Attempt to parse it using SQLGlot's BigQuery dialect.
##### Observe the failure in parsing the query.
##### Proposed Solution
##### To address this limitation:

#### Extend SQLGlot's BigQuery dialect to include parsing rules for ML.FEATURES_AT_TIME.
#### Define a custom expression in SQLGlot for this function.
#### Update the generator to ensure proper SQL generation from the AST.

Hints from issue comments:
@VaggelisD: Hey @YuvrajSoni-Ksolves,

Thank you for reporting this. Sadly some functions such as `FEATURES_AT_TIME()` require more nuanced handling; if they're not manually added to SQLGlot, their parsing can fail.

I went ahead and implemented `FEATURES_AT_TIME()` but our bandwidth is limited so we'd not be able to support all other ML functions right away. 

If possible, consider taking a look at my approach and contributing any other problematic function you come across, we accept well-crafted PRs.

Solution patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 65ff1b224a..b7802f890a 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -534,6 +534,7 @@ class Parser(parser.Parser):
             **parser.Parser.FUNCTION_PARSERS,
             "ARRAY": lambda self: self.expression(exp.Array, expressions=[self._parse_statement()]),
             "MAKE_INTERVAL": lambda self: self._parse_make_interval(),
+            "FEATURES_AT_TIME": lambda self: self._parse_features_at_time(),
         }
         FUNCTION_PARSERS.pop("TRIM")
 
@@ -764,7 +765,7 @@ def _parse_unnest(self, with_alias: bool = True) -> t.Optional[exp.Unnest]:
 
             return unnest
 
-        def _parse_make_interval(self):
+        def _parse_make_interval(self) -> exp.MakeInterval:
             expr = exp.MakeInterval()
 
             for arg_key in expr.arg_types:
@@ -784,6 +785,23 @@ def _parse_make_interval(self):
 
             return expr
 
+        def _parse_features_at_time(self) -> exp.FeaturesAtTime:
+            expr = self.expression(
+                exp.FeaturesAtTime,
+                this=(self._match(TokenType.TABLE) and self._parse_table())
+                or self._parse_select(nested=True),
+            )
+
+            while self._match(TokenType.COMMA):
+                arg = self._parse_lambda()
+
+                # Get the LHS of the Kwarg and set the arg to that value, e.g
+                # "num_rows => 1" sets the expr's `num_rows` arg
+                if arg:
+                    expr.set(arg.this.name, arg)
+
+            return expr
+
     class Generator(generator.Generator):
         INTERVAL_ALLOWS_PLURAL_FORM = False
         JOIN_HINTS = False
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index a69e8362f9..ebbc94d9a2 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -5772,6 +5772,10 @@ class FromBase64(Func):
     pass
 
 
+class FeaturesAtTime(Func):
+    arg_types = {"this": True, "time": False, "num_rows": False, "ignore_feature_nulls": False}
+
+
 class ToBase64(Func):
     pass
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index ae9dc35e49..83222cdccc 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -4543,3 +4543,16 @@ def attachoption_sql(self, expression: exp.AttachOption) -> str:
         value = self.sql(expression, "expression")
         value = f" {value}" if value else ""
         return f"{this}{value}"
+
+    def featuresattime_sql(self, expression: exp.FeaturesAtTime) -> str:
+        this_sql = self.sql(expression, "this")
+        if isinstance(expression.this, exp.Table):
+            this_sql = f"TABLE {this_sql}"
+
+        return self.func(
+            "FEATURES_AT_TIME",
+            this_sql,
+            expression.args.get("time"),
+            expression.args.get("num_rows"),
+            expression.args.get("ignore_feature_nulls"),
+        )


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index f88138554f..f80c3ac5b0 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -1640,6 +1640,11 @@ def test_bigquery(self):
             },
         )
 
+        self.validate_identity(
+            "SELECT * FROM ML.FEATURES_AT_TIME(TABLE mydataset.feature_table, time => '2022-06-11 10:00:00+00', num_rows => 1, ignore_feature_nulls => TRUE)"
+        )
+        self.validate_identity("SELECT * FROM ML.FEATURES_AT_TIME((SELECT 1), num_rows => 1)")
+
     def test_errors(self):
         with self.assertRaises(TokenError):
             transpile("'\\'", read="bigquery")