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.

Discussion:

Patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 65ff1b22..b7802f89 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -534,6 +534,7 @@ class BigQuery(Dialect):
             **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 @@ class BigQuery(Dialect):
 
             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 @@ class BigQuery(Dialect):
 
             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 a69e8362..ebbc94d9 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 ae9dc35e..83222cdc 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -4543,3 +4543,16 @@ class Generator(metaclass=_Generator):
         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 f8813855..f80c3ac5 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -1640,6 +1640,11 @@ WHERE
             },
         )
 
+        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")