diff --git a/vertexai/resources/preview/feature_store/_offline_store_impl.py b/vertexai/resources/preview/feature_store/_offline_store_impl.py index 026b74e79e..3e044dc7a7 100644 --- a/vertexai/resources/preview/feature_store/_offline_store_impl.py +++ b/vertexai/resources/preview/feature_store/_offline_store_impl.py @@ -86,6 +86,11 @@ def qualified_timestamp_column(self) -> str: """Returns name qualified timestamp column e.g. `name.feature_timestamp`.""" return f"{self.qualifying_name}.{self.timestamp_column}" + @property + def comma_separated_row_num_columns(self) -> str: + """Returns column names for the order by used in the initial CTE as a salting strategy.""" + return ",".join(self.data_columns + [self.timestamp_column]) + def _generate_eid_check(entity_data: DataSource, feature: DataSource): """Generate equality check for entity columns of feature against matching columns in entity_data.""" @@ -118,7 +123,9 @@ def _generate_eid_check(entity_data: DataSource, feature: DataSource): {{ textwrap.indent(entity_data.sql, ' ' * 4) }} ), {{ entity_data.qualifying_name }} AS ( - SELECT *, ROW_NUMBER() OVER() AS row_num, + SELECT *, ROW_NUMBER() OVER( + ORDER BY {{ entity_data.comma_separated_row_num_columns }} + ) AS row_num, FROM entity_df_without_row_num ), diff --git a/vertexai/resources/preview/feature_store/offline_store.py b/vertexai/resources/preview/feature_store/offline_store.py index 3a785711e4..4afccdf725 100644 --- a/vertexai/resources/preview/feature_store/offline_store.py +++ b/vertexai/resources/preview/feature_store/offline_store.py @@ -214,7 +214,7 @@ def fetch_historical_feature_values( session_options = bigframes.BigQueryOptions( credentials=credentials, project=project, - location=location, + location="US", application_name=application_name, ) session = bigframes.connect(session_options)