import contextlib

import json

import os

import tempfile

import uuid

from datetime import date, datetime, timedelta, timezone

from pathlib import Path

from typing import (

Any,

Callable,

ContextManager,

Dict,

Iterator,

KeysView,

List,

Literal,

Optional,

Tuple,

Union,

cast,

)

import numpy as np

import pandas as pd

import pyarrow

import pyarrow.parquet

from pydantic import StrictStr, field_validator

from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed

from feast import flags_helper

from feast.credentials import get_connection_config_override

from feast.data_source import DataSource

from feast.errors import (

BigQueryJobCancelled,

BigQueryJobStillRunning,

EntityDFNotDateTime,

EntitySQLEmptyResults,

FeastProviderLoginError,

InvalidEntityType,

)

from feast.feature_logging import LoggingConfig, LoggingSource

from feast.feature_view import DUMMY_ENTITY_ID, DUMMY_ENTITY_VAL, FeatureView

from feast.infra.offline_stores import offline_utils

from feast.infra.offline_stores.offline_store import (

OfflineStore,

RetrievalJob,

RetrievalMetadata,

)

from feast.infra.registry.base_registry import BaseRegistry

from feast.monitoring.monitoring_utils import (

MON_TABLE_FEATURE,

MON_TABLE_FEATURE_SERVICE,

MON_TABLE_FEATURE_VIEW,

MON_TABLE_JOB,

empty_categorical_metric,

empty_numeric_metric,

monitoring_table_meta,

normalize_monitoring_row,

opt_float,

)

from feast.on_demand_feature_view import OnDemandFeatureView

from feast.repo_config import FeastConfigBaseModel, RepoConfig

from feast.saved_dataset import SavedDatasetStorage

from feast.utils import _utc_now, compute_non_entity_date_range, get_user_agent

from .bigquery_source import (

BigQueryLoggingDestination,

BigQuerySource,

SavedDatasetBigQueryStorage,

)

from .offline_utils import get_timestamp_filter_sql

try:

from google.api_core import client_info as http_client_info

from google.api_core.exceptions import NotFound

from google.auth.exceptions import DefaultCredentialsError

from google.cloud import bigquery

from google.cloud.bigquery import Client, SchemaField, Table

from google.cloud.storage import Client as StorageClient

except ImportError as e:

from feast.errors import FeastExtrasDependencyImportError

raise FeastExtrasDependencyImportError("gcp", str(e))

try:

from google.cloud.bigquery._pyarrow_helpers import _ARROW_SCALAR_IDS_TO_BQ

except ImportError:

try:

from google.cloud.bigquery._pandas_helpers import ( # type: ignore

ARROW_SCALAR_IDS_TO_BQ as _ARROW_SCALAR_IDS_TO_BQ,

)

except ImportError as e:

raise FeastExtrasDependencyImportError("gcp", str(e))

def get_http_client_info():

return http_client_info.ClientInfo(user_agent=get_user_agent())

class BigQueryOfflineStoreConfig(FeastConfigBaseModel):

"""Offline store config for GCP BigQuery"""

type: Literal["bigquery"] = "bigquery"

""" Offline store type selector"""

dataset: StrictStr = "feast"

""" (optional) BigQuery Dataset name for temporary tables """

project_id: Optional[StrictStr] = None

""" (optional) GCP project name used for the BigQuery offline store """

billing_project_id: Optional[StrictStr] = None

""" (optional) GCP project name used to run the bigquery jobs at """

location: Optional[StrictStr] = None

""" (optional) GCP location name used for the BigQuery offline store.

Examples of location names include ``US``, ``EU``, ``us-central1``, ``us-west4``.

If a location is not specified, the location defaults to the ``US`` multi-regional location.

For more information on BigQuery data locations see: https://cloud.google.com/bigquery/docs/locations

"""

gcs_staging_location: Optional[str] = None

""" (optional) GCS location used for offloading BigQuery results as parquet files."""

table_create_disposition: Literal["CREATE_NEVER", "CREATE_IF_NEEDED"] = (

"CREATE_IF_NEEDED"

)

""" (optional) Specifies whether the job is allowed to create new tables. The default value is CREATE_IF_NEEDED.

Custom constraint for table_create_disposition. To understand more, see:

https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationLoad.FIELDS.create_disposition

"""

@field_validator("billing_project_id")

def project_id_exists(cls, v, values, **kwargs):

if v and not values.data["project_id"]:

raise ValueError(

"please specify project_id if billing_project_id is specified"

)

return v

class BigQueryOfflineStore(OfflineStore):

supports_filter_by_created_timestamp = True

@staticmethod

def pull_latest_from_table_or_query(

config: RepoConfig,

data_source: DataSource,

join_key_columns: List[str],

feature_name_columns: List[str],

timestamp_field: str,

created_timestamp_column: Optional[str],

start_date: datetime,

end_date: datetime,

) -> RetrievalJob:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

assert isinstance(data_source, BigQuerySource)

from_expression = data_source.get_table_query_string()

partition_by_join_key_string = ", ".join(

BigQueryOfflineStore._escape_query_columns(join_key_columns)

)

if partition_by_join_key_string != "":

partition_by_join_key_string = (

"PARTITION BY " + partition_by_join_key_string

)

timestamps = [timestamp_field]

if created_timestamp_column:

timestamps.append(created_timestamp_column)

timestamp_desc_string = " DESC, ".join(timestamps) + " DESC"

field_string = ", ".join(

BigQueryOfflineStore._escape_query_columns(join_key_columns)

+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)

+ timestamps

)

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

cast_style: Literal["date_func", "timestamp_func"] = (

"date_func"

if data_source.timestamp_field_type == "DATE"

else "timestamp_func"

)

timestamp_filter = get_timestamp_filter_sql(

start_date,

end_date,

timestamp_field,

date_partition_column=data_source.date_partition_column,

quote_fields=False,

cast_style=cast_style,

)

query = f"""

SELECT

{field_string}

{f", {repr(DUMMY_ENTITY_VAL)} AS {DUMMY_ENTITY_ID}" if not join_key_columns else ""}

FROM (

SELECT {field_string},

ROW_NUMBER() OVER({partition_by_join_key_string} ORDER BY {timestamp_desc_string}) AS _feast_row

FROM {from_expression}

WHERE {timestamp_filter}

)

WHERE _feast_row = 1

"""

# When materializing a single feature view, we don't need full feature names. On demand transforms aren't materialized

return BigQueryRetrievalJob(

query=query,

client=client,

config=config,

full_feature_names=False,

)

@staticmethod

def pull_all_from_table_or_query(

config: RepoConfig,

data_source: DataSource,

join_key_columns: List[str],

feature_name_columns: List[str],

timestamp_field: str,

created_timestamp_column: Optional[str] = None,

start_date: Optional[datetime] = None,

end_date: Optional[datetime] = None,

) -> RetrievalJob:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

assert isinstance(data_source, BigQuerySource)

from_expression = data_source.get_table_query_string()

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

timestamp_fields = [timestamp_field]

if created_timestamp_column:

timestamp_fields.append(created_timestamp_column)

field_string = ", ".join(

BigQueryOfflineStore._escape_query_columns(join_key_columns)

+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)

+ timestamp_fields

)

cast_style: Literal["date_func", "timestamp_func"] = (

"date_func"

if data_source.timestamp_field_type == "DATE"

else "timestamp_func"

)

timestamp_filter = get_timestamp_filter_sql(

start_date,

end_date,

timestamp_field,

date_partition_column=data_source.date_partition_column,

quote_fields=False,

cast_style=cast_style,

)

query = f"""

SELECT {field_string}

FROM {from_expression}

WHERE {timestamp_filter}

"""

return BigQueryRetrievalJob(

query=query,

client=client,

config=config,

full_feature_names=False,

)

@staticmethod

def get_historical_features(

config: RepoConfig,

feature_views: List[FeatureView],

feature_refs: List[str],

entity_df: Optional[Union[pd.DataFrame, str]],

registry: BaseRegistry,

project: str,

full_feature_names: bool = False,

filter_by_created_timestamp: bool = False,

**kwargs: Any,

) -> RetrievalJob:

# TODO: Add entity_df validation in order to fail before interacting with BigQuery

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

for fv in feature_views:

assert isinstance(fv.batch_source, BigQuerySource)

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

if config.offline_store.billing_project_id:

dataset_project = str(config.offline_store.project_id)

else:

dataset_project = client.project

table_reference = _get_table_reference_for_new_entity(

client,

dataset_project,

config.offline_store.dataset,

config.offline_store.location,

config.offline_store.table_create_disposition,

)

# Non-entity mode: create a left temporary table from entity keys - any entity key having an event in the time window

non_entity_mode = entity_df is None

if non_entity_mode:

start_date, end_date = compute_non_entity_date_range(

feature_views,

start_date=kwargs.get("start_date"),

end_date=kwargs.get("end_date"),

)

entity_df_event_timestamp_range = (start_date, end_date)

# Pre-compute query contexts to collect entity column names per feature view.

fv_query_contexts_pre = offline_utils.get_feature_view_query_context(

feature_refs,

feature_views,

registry,

project,

entity_df_event_timestamp_range,

)

all_entities = offline_utils.gather_all_entities(fv_query_contexts_pre)

event_timestamp_col = "entity_ts"

entity_schema_keys: KeysView[str] = cast(

KeysView[str],

{k: None for k in (all_entities + [event_timestamp_col])}.keys(),

)

entity_schema = None

else:

entity_schema = _get_entity_schema(

client=client,

entity_df=entity_df,

)

event_timestamp_col = offline_utils.infer_event_timestamp_from_entity_df(

entity_schema

)

entity_df_event_timestamp_range = _get_entity_df_event_timestamp_range(

entity_df,

event_timestamp_col,

client,

)

entity_schema_keys = entity_schema.keys()

all_entities = []

start_date = entity_df_event_timestamp_range[0]

end_date = entity_df_event_timestamp_range[1]

fv_query_contexts_pre = offline_utils.get_feature_view_query_context(

feature_refs,

feature_views,

registry,

project,

entity_df_event_timestamp_range,

)

@contextlib.contextmanager

def query_generator() -> Iterator[str]:

if non_entity_mode:

_bq_create_entity_union_table(

client=client,

table_name=table_reference,

feature_views=feature_views,

fv_query_contexts=fv_query_contexts_pre,

start_date=start_date,

end_date=end_date,

all_entities=all_entities,

event_timestamp_col=event_timestamp_col,

)

else:

_upload_entity_df(

client=client,

table_name=table_reference,

entity_df=entity_df,

)

expected_join_keys = offline_utils.get_expected_join_keys(

project, feature_views, registry

)

assert entity_schema is not None

offline_utils.assert_expected_columns_in_entity_df(

entity_schema, expected_join_keys, event_timestamp_col

)

# Generate the BigQuery SQL query from the query context

query = offline_utils.build_point_in_time_query(

feature_view_query_contexts=fv_query_contexts_pre,

left_table_query_string=table_reference,

entity_df_event_timestamp_col=event_timestamp_col,

entity_df_columns=entity_schema_keys,

query_template=MULTIPLE_FEATURE_VIEW_POINT_IN_TIME_JOIN,

full_feature_names=full_feature_names,

filter_by_created_timestamp=filter_by_created_timestamp,

)

try:

yield query

finally:

# Asynchronously clean up the uploaded BigQuery table, which will expire

# if cleanup fails

client.delete_table(table=table_reference, not_found_ok=True)

return BigQueryRetrievalJob(

query=query_generator,

client=client,

config=config,

full_feature_names=full_feature_names,

on_demand_feature_views=OnDemandFeatureView.get_requested_odfvs(

feature_refs, project, registry

),

metadata=RetrievalMetadata(

features=feature_refs,

keys=list(set(entity_schema_keys) - {event_timestamp_col}),

min_event_timestamp=entity_df_event_timestamp_range[0],

max_event_timestamp=entity_df_event_timestamp_range[1],

),

)

@staticmethod

def write_logged_features(

config: RepoConfig,

data: Union[pyarrow.Table, Path],

source: LoggingSource,

logging_config: LoggingConfig,

registry: BaseRegistry,

):

destination = logging_config.destination

assert isinstance(destination, BigQueryLoggingDestination)

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

job_config = bigquery.LoadJobConfig(

source_format=bigquery.SourceFormat.PARQUET,

schema=arrow_schema_to_bq_schema(source.get_schema(registry)),

create_disposition=config.offline_store.table_create_disposition,

time_partitioning=bigquery.TimePartitioning(

type_=bigquery.TimePartitioningType.DAY,

field=source.get_log_timestamp_column(),

),

)

if isinstance(data, Path):

for file in data.iterdir():

with file.open("rb") as f:

client.load_table_from_file(

file_obj=f,

destination=destination.table,

job_config=job_config,

).result()

return

with tempfile.TemporaryFile() as parquet_temp_file:

# In Pyarrow v13.0, the parquet version was upgraded to v2.6 from v2.4.

# Set the coerce_timestamps to "us"(microseconds) for backward compatibility.

pyarrow.parquet.write_table(

table=data,

where=parquet_temp_file,

coerce_timestamps="us",

allow_truncated_timestamps=True,

)

parquet_temp_file.seek(0)

client.load_table_from_file(

file_obj=parquet_temp_file,

destination=destination.table,

job_config=job_config,

).result()

@staticmethod

def offline_write_batch(

config: RepoConfig,

feature_view: FeatureView,

table: pyarrow.Table,

progress: Optional[Callable[[int], Any]],

):

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

assert isinstance(feature_view.batch_source, BigQuerySource)

pa_schema, column_names = offline_utils.get_pyarrow_schema_from_batch_source(

config, feature_view.batch_source, timestamp_unit="ns"

)

if column_names != table.column_names:

raise ValueError(

f"The input pyarrow table has schema {table.schema} with the incorrect columns {table.column_names}. "

f"The schema is expected to be {pa_schema} with the columns (in this exact order) to be {column_names}."

)

if table.schema != pa_schema:

table = offline_utils.cast_arrow_table_to_schema(table, pa_schema)

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

parquet_options = bigquery.ParquetOptions()

parquet_options.enable_list_inference = True

job_config = bigquery.LoadJobConfig(

source_format=bigquery.SourceFormat.PARQUET,

schema=arrow_schema_to_bq_schema(pa_schema),

create_disposition=config.offline_store.table_create_disposition,

write_disposition="WRITE_APPEND", # Default but included for clarity

parquet_options=parquet_options,

)

with tempfile.TemporaryFile() as parquet_temp_file:

# In Pyarrow v13.0, the parquet version was upgraded to v2.6 from v2.4.

# Set the coerce_timestamps to "us"(microseconds) for backward compatibility.

pyarrow.parquet.write_table(

table=table,

where=parquet_temp_file,

coerce_timestamps="us",

allow_truncated_timestamps=True,

)

parquet_temp_file.seek(0)

client.load_table_from_file(

file_obj=parquet_temp_file,

destination=feature_view.batch_source.table,

job_config=job_config,

).result()

@staticmethod

def _escape_query_columns(columns: List[str]) -> List[str]:

return [f"`{x}`" for x in columns]

@staticmethod

def compute_monitoring_metrics(

config: RepoConfig,

data_source: DataSource,

feature_columns: List[Tuple[str, str]],

timestamp_field: str,

start_date: Optional[datetime] = None,

end_date: Optional[datetime] = None,

histogram_bins: int = 20,

top_n: int = 10,

) -> List[Dict[str, Any]]:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

assert isinstance(data_source, BigQuerySource)

return _bq_compute_monitoring_metrics(

config,

data_source,

feature_columns,

timestamp_field,

start_date=start_date,

end_date=end_date,

histogram_bins=histogram_bins,

top_n=top_n,

)

@staticmethod

def get_monitoring_max_timestamp(

config: RepoConfig,

data_source: DataSource,

timestamp_field: str,

) -> Optional[datetime]:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

assert isinstance(data_source, BigQuerySource)

return _bq_get_monitoring_max_timestamp(config, data_source, timestamp_field)

@staticmethod

def ensure_monitoring_tables(config: RepoConfig) -> None:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

_bq_ensure_monitoring_tables(config)

@staticmethod

def save_monitoring_metrics(

config: RepoConfig,

metric_type: str,

metrics: List[Dict[str, Any]],

) -> None:

if not metrics:

return

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

_bq_save_monitoring_metrics(config, metric_type, metrics)

@staticmethod

def query_monitoring_metrics(

config: RepoConfig,

project: str,

metric_type: str,

filters: Optional[Dict[str, Any]] = None,

start_date: Optional[date] = None,

end_date: Optional[date] = None,

) -> List[Dict[str, Any]]:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

return _bq_query_monitoring_metrics(

config, project, metric_type, filters, start_date, end_date

)

@staticmethod

def clear_monitoring_baseline(

config: RepoConfig,

project: str,

feature_view_name: Optional[str] = None,

feature_name: Optional[str] = None,

data_source_type: Optional[str] = None,

) -> None:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

_bq_clear_monitoring_baseline(

config, project, feature_view_name, feature_name, data_source_type

)

def _bq_create_entity_union_table(

client: "Client",

table_name: str,

feature_views: List[FeatureView],

fv_query_contexts: List[offline_utils.FeatureViewQueryContext],

start_date: datetime,

end_date: datetime,

all_entities: List[str],

event_timestamp_col: str,

) -> None:

"""

Creates a BigQuery temp table containing the UNION DISTINCT of entity keys observed

across all feature views in [start_date, end_date], plus a stable as-of timestamp

column set to end_date. Used as the left table for PIT joins in non-entity mode.

"""

start_str = start_date.strftime("%Y-%m-%dT%H:%M:%S")

end_str = end_date.strftime("%Y-%m-%dT%H:%M:%S")

per_view_selects: List[str] = []

for fv, ctx in zip(feature_views, fv_query_contexts):

assert isinstance(fv.batch_source, BigQuerySource)

from_expression = fv.batch_source.get_table_query_string()

timestamp_field = ctx.timestamp_field

ctx_entities_set = set(ctx.entities)

select_entities: List[str] = []

for col in all_entities:

if col in ctx_entities_set:

select_entities.append(f"`{col}`")

else:

select_entities.append(f"CAST(NULL AS STRING) AS `{col}`")

per_view_selects.append(

f"SELECT DISTINCT {', '.join(select_entities)} "

f"FROM {from_expression} "

f"WHERE `{timestamp_field}` BETWEEN TIMESTAMP('{start_str}') AND TIMESTAMP('{end_str}')"

)

union_query = "\nUNION DISTINCT\n".join(per_view_selects)

entity_cols = (

", ".join(f"`{e}`" for e in all_entities) if all_entities else "TRUE AS _dummy"

)

create_sql = (

f"CREATE TABLE `{table_name}` AS "

f"SELECT {entity_cols}, TIMESTAMP('{end_str}') AS `{event_timestamp_col}` "

f"FROM ({union_query}) AS _entity_union"

)

block_until_done(client, client.query(create_sql))

table = client.get_table(table_name)

table.expires = _utc_now() + timedelta(minutes=30)

client.update_table(table, ["expires"])

# ------------------------------------------------------------------ #

# BigQuery monitoring metrics (native)

# ------------------------------------------------------------------ #

def _bq_monitoring_table_fqn(config: RepoConfig, table_name: str) -> str:

assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)

project_id = config.offline_store.project_id

if not project_id:

client = _get_bigquery_client(

project=config.offline_store.billing_project_id,

location=config.offline_store.location,

)

project_id = client.project

return f"`{project_id}.{config.offline_store.dataset}.{table_name}`"

def _bq_scalar_param_type(column: str) -> str:

if column == "is_baseline":

return "BOOL"

if column == "metric_date":

return "DATE"

if column in ("computed_at", "max_event_timestamp"):

return "TIMESTAMP"

if column in {

"row_count",

"null_count",

"total_row_count",

"total_features",

"features_with_nulls",

"total_feature_views",

}:

return "INT64"

if column in {

"null_rate",

"mean",

"stddev",

"min_val",

"max_val",

"p50",

"p75",

"p90",

"p95",

"p99",

"avg_null_rate",

"max_null_rate",

}:

return "FLOAT64"

return "STRING"

def _bq_merge_row(

config: RepoConfig,

table_fqn: str,

columns: List[str],

pk_columns: List[str],

row: Dict[str, Any],

) -> None:

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

non_pk = [c for c in columns if c not in pk_columns]

params: List[Any] = []

using_parts: List[str] = []

on_parts: List[str] = []

merge_idx = 0

for col in columns:

p = f"p{merge_idx}"

merge_idx += 1

val = row.get(col)

if col == "histogram" and val is not None and not isinstance(val, str):

val = json.dumps(val)

param_type = _bq_scalar_param_type(col)

params.append(bigquery.ScalarQueryParameter(p, param_type, val))

using_parts.append(f"@{p} AS {col}")

on_parts = [f"T.{col} = S.{col}" for col in pk_columns]

update_set = ", ".join(f"{c} = S.{c}" for c in non_pk)

merge_sql = f"""

MERGE {table_fqn} T

USING (SELECT {", ".join(using_parts)}) S

ON {" AND ".join(on_parts)}

WHEN MATCHED THEN UPDATE SET {update_set}

WHEN NOT MATCHED THEN INSERT ({", ".join(columns)}) VALUES ({", ".join(f"S.{c}" for c in columns)})

"""

job_config = bigquery.QueryJobConfig(query_parameters=params)

client.query(merge_sql, job_config=job_config).result()

def _bq_save_monitoring_metrics(

config: RepoConfig,

metric_type: str,

metrics: List[Dict[str, Any]],

) -> None:

table_short, columns, pk_columns = monitoring_table_meta(metric_type)

table_fqn = _bq_monitoring_table_fqn(config, table_short)

for row in metrics:

_bq_merge_row(config, table_fqn, columns, pk_columns, row)

def _bq_query_monitoring_metrics(

config: RepoConfig,

project: str,

metric_type: str,

filters: Optional[Dict[str, Any]] = None,

start_date: Optional[date] = None,

end_date: Optional[date] = None,

) -> List[Dict[str, Any]]:

table_short, columns, _ = monitoring_table_meta(metric_type)

table_fqn = _bq_monitoring_table_fqn(config, table_short)

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

params: List[Any] = []

conditions: List[str] = []

if project:

params.append(

bigquery.ScalarQueryParameter("project", "STRING", project),

)

conditions.append("project_id = @project")

if filters:

for key, value in filters.items():

if value is not None:

conditions.append(f"`{key}` = @{key}")

params.append(

bigquery.ScalarQueryParameter(

key, _bq_scalar_param_type(key), value

)

)

if start_date:

conditions.append("metric_date >= @start_date")

params.append(bigquery.ScalarQueryParameter("start_date", "DATE", start_date))

if end_date:

conditions.append("metric_date <= @end_date")

params.append(bigquery.ScalarQueryParameter("end_date", "DATE", end_date))

col_list = ", ".join(f"`{c}`" for c in columns)

where_sql = " AND ".join(conditions) if conditions else "TRUE"

order_col = "metric_date" if "metric_date" in columns else "job_id"

sql = f"SELECT {col_list} FROM {table_fqn} WHERE {where_sql} ORDER BY `{order_col}` ASC"

job_config = bigquery.QueryJobConfig(query_parameters=params)

job = client.query(sql, job_config=job_config)

job.result()

results: List[Dict[str, Any]] = []

for r in job:

record = {columns[i]: r[i] for i in range(len(columns))}

results.append(normalize_monitoring_row(record))

return results

def _bq_clear_monitoring_baseline(

config: RepoConfig,

project: str,

feature_view_name: Optional[str] = None,

feature_name: Optional[str] = None,

data_source_type: Optional[str] = None,

) -> None:

table_fqn = _bq_monitoring_table_fqn(config, MON_TABLE_FEATURE)

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

params: List[Any] = [

bigquery.ScalarQueryParameter("project", "STRING", project),

]

conditions = ["project_id = @project", "is_baseline = TRUE"]

if feature_view_name:

conditions.append("feature_view_name = @feature_view_name")

params.append(

bigquery.ScalarQueryParameter(

"feature_view_name", "STRING", feature_view_name

)

)

if feature_name:

conditions.append("feature_name = @feature_name")

params.append(

bigquery.ScalarQueryParameter("feature_name", "STRING", feature_name)

)

if data_source_type:

conditions.append("data_source_type = @data_source_type")

params.append(

bigquery.ScalarQueryParameter(

"data_source_type", "STRING", data_source_type

)

)

sql = f"UPDATE {table_fqn} SET is_baseline = FALSE WHERE {' AND '.join(conditions)}"

job_config = bigquery.QueryJobConfig(query_parameters=params)

client.query(sql, job_config=job_config).result()

def _bq_ensure_monitoring_tables(config: RepoConfig) -> None:

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

ds = config.offline_store.dataset

proj = config.offline_store.project_id or client.project

feature_ddl = f"""

CREATE TABLE IF NOT EXISTS `{proj}.{ds}.{MON_TABLE_FEATURE}` (

project_id STRING NOT NULL,

feature_view_name STRING NOT NULL,

feature_name STRING NOT NULL,

metric_date DATE NOT NULL,

granularity STRING NOT NULL,

data_source_type STRING NOT NULL,

computed_at TIMESTAMP NOT NULL,

max_event_timestamp TIMESTAMP,

is_baseline BOOL NOT NULL,

feature_type STRING NOT NULL,

row_count INT64,

null_count INT64,

null_rate FLOAT64,

mean FLOAT64,

stddev FLOAT64,

min_val FLOAT64,

max_val FLOAT64,

p50 FLOAT64,

p75 FLOAT64,

p90 FLOAT64,

p95 FLOAT64,

p99 FLOAT64,

histogram STRING

)

PRIMARY KEY (project_id, feature_view_name, feature_name, metric_date, granularity, data_source_type) NOT ENFORCED

"""

view_ddl = f"""

CREATE TABLE IF NOT EXISTS `{proj}.{ds}.{MON_TABLE_FEATURE_VIEW}` (

project_id STRING NOT NULL,

feature_view_name STRING NOT NULL,

metric_date DATE NOT NULL,

granularity STRING NOT NULL,

data_source_type STRING NOT NULL,

computed_at TIMESTAMP NOT NULL,

max_event_timestamp TIMESTAMP,

is_baseline BOOL NOT NULL,

total_row_count INT64,

total_features INT64,

features_with_nulls INT64,

avg_null_rate FLOAT64,

max_null_rate FLOAT64

)

PRIMARY KEY (project_id, feature_view_name, metric_date, granularity, data_source_type) NOT ENFORCED

"""

service_ddl = f"""

CREATE TABLE IF NOT EXISTS `{proj}.{ds}.{MON_TABLE_FEATURE_SERVICE}` (

project_id STRING NOT NULL,

feature_service_name STRING NOT NULL,

metric_date DATE NOT NULL,

granularity STRING NOT NULL,

data_source_type STRING NOT NULL,

computed_at TIMESTAMP NOT NULL,

max_event_timestamp TIMESTAMP,

is_baseline BOOL NOT NULL,

total_feature_views INT64,

total_features INT64,

avg_null_rate FLOAT64,

max_null_rate FLOAT64

)

PRIMARY KEY (project_id, feature_service_name, metric_date, granularity, data_source_type) NOT ENFORCED

"""

job_ddl = f"""

CREATE TABLE IF NOT EXISTS `{proj}.{ds}.{MON_TABLE_JOB}` (

job_id STRING NOT NULL,

project_id STRING NOT NULL,

feature_view_name STRING,

job_type STRING NOT NULL,

status STRING NOT NULL,

parameters STRING,

metric_date DATE NOT NULL,

started_at TIMESTAMP,

completed_at TIMESTAMP,

error_message STRING,

result_summary STRING

)

PRIMARY KEY (job_id) NOT ENFORCED

"""

for ddl in (feature_ddl, view_ddl, service_ddl, job_ddl):

client.query(ddl).result()

for tbl in (MON_TABLE_FEATURE, MON_TABLE_FEATURE_VIEW, MON_TABLE_FEATURE_SERVICE):

client.query(

f"ALTER TABLE `{proj}.{ds}.{tbl}` "

"ADD COLUMN IF NOT EXISTS max_event_timestamp TIMESTAMP"

).result()

def _bq_get_monitoring_max_timestamp(

config: RepoConfig,

data_source: BigQuerySource,

timestamp_field: str,

) -> Optional[datetime]:

from_expression = data_source.get_table_query_string()

ts_col = f"`{timestamp_field}`"

sql = f"SELECT MAX({ts_col}) AS _max_ts FROM {from_expression}"

project_id = (

config.offline_store.billing_project_id or config.offline_store.project_id

)

client = _get_bigquery_client(

project=project_id,

location=config.offline_store.location,

)

job = client.query(sql)

job.result()

rows = list(job)

if not rows or rows[0][0] is None:

return None

val = rows[0][0]

if isinstance(val, datetime):

return val if val.tzinfo else val.replace(tzinfo=timezone.utc)

if isinstance(val, date):

return datetime.combine(val, datetime.min.time(), tzinfo=timezone.utc)

return val # type: ignore[no-any-return]

def _bq_numeric_histogram(