diff --git a/infra/scripts/test-end-to-end-gcp.sh b/infra/scripts/test-end-to-end-gcp.sh index 20f628bfed7..bc3b8352870 100755 --- a/infra/scripts/test-end-to-end-gcp.sh +++ b/infra/scripts/test-end-to-end-gcp.sh @@ -10,7 +10,8 @@ python -m pip install --upgrade pip setuptools wheel make install-python python -m pip install -qr tests/requirements.txt -su -p postgres -c "PATH=$PATH HOME=/tmp pytest tests/e2e/ \ +su -p postgres -c "PATH=$PATH HOME=/tmp pytest -v tests/e2e/ \ --feast-version develop --env=gcloud --dataproc-cluster-name feast-e2e \ --dataproc-project kf-feast --dataproc-region us-central1 \ - --redis-url 10.128.0.105:6379 --redis-cluster --kafka-brokers 10.128.0.103:9094" + --redis-url 10.128.0.105:6379 --redis-cluster --kafka-brokers 10.128.0.103:9094 \ + --bq-project kf-feast" diff --git a/infra/scripts/test-end-to-end.sh b/infra/scripts/test-end-to-end.sh index 127ea474337..60f9c33a140 100755 --- a/infra/scripts/test-end-to-end.sh +++ b/infra/scripts/test-end-to-end.sh @@ -7,4 +7,4 @@ python -m pip install --upgrade pip setuptools wheel make install-python python -m pip install -qr tests/requirements.txt -su -p postgres -c "PATH=$PATH HOME=/tmp pytest tests/e2e/ --feast-version develop" \ No newline at end of file +su -p postgres -c "PATH=$PATH HOME=/tmp pytest -v tests/e2e/ --feast-version develop" \ No newline at end of file diff --git a/infra/scripts/test-integration.sh b/infra/scripts/test-integration.sh old mode 100644 new mode 100755 diff --git a/protos/feast/core/JobService.proto b/protos/feast/core/JobService.proto index 9788121cc23..13e37b7ffe2 100644 --- a/protos/feast/core/JobService.proto +++ b/protos/feast/core/JobService.proto @@ -128,6 +128,9 @@ message GetHistoricalFeaturesRequest { // Export to AWS S3 - s3://path/to/features // Export to GCP GCS - gs://path/to/features string output_location = 4; + + // Specify format name for output, eg. parquet + string output_format = 5; } message GetHistoricalFeaturesResponse { diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py index 9a959a8afea..6c57183f8a4 100644 --- a/sdk/python/feast/client.py +++ b/sdk/python/feast/client.py @@ -15,12 +15,10 @@ import multiprocessing import os import shutil -import tempfile import uuid from datetime import datetime from itertools import groupby from typing import Any, Dict, List, Optional, Union -from urllib.parse import urlparse import grpc import pandas as pd @@ -101,7 +99,11 @@ GetOnlineFeaturesRequestV2, ) from feast.serving.ServingService_pb2_grpc import ServingServiceStub -from feast.staging.storage_client import get_staging_client +from feast.staging.entities import ( + stage_entities_to_bq, + stage_entities_to_fs, + table_reference_from_string, +) _logger = logging.getLogger(__name__) @@ -855,6 +857,7 @@ def get_online_features( entity_rows=_infer_online_entity_rows(entity_rows), project=project if project is not None else self.project, ), + timeout=self._config.getint(CONFIG_GRPC_CONNECTION_TIMEOUT_DEFAULT_KEY), metadata=self._get_grpc_metadata(), ) except grpc.RpcError as e: @@ -879,8 +882,11 @@ def get_historical_features( "feature_table:feature" where "feature_table" & "feature" refer to the feature and feature table names respectively. entity_source (Union[pd.DataFrame, FileSource, BigQuerySource]): Source for the entity rows. - If entity_source is a Panda DataFrame, the dataframe will be exported to the staging - location as parquet file. It is also assumed that the column event_timestamp is present + If entity_source is a Panda DataFrame, the dataframe will be staged + to become accessible by spark workers. + If one of feature tables' source is in BigQuery - entities will be upload to BQ. + Otherwise to remote file storage (derived from configured staging location). + It is also assumed that the column event_timestamp is present in the dataframe, and is of type datetime without timezone information. The user needs to make sure that the source (or staging location, if entity_source is @@ -916,25 +922,27 @@ def get_historical_features( str(uuid.uuid4()), ) output_format = self._config.get(CONFIG_SPARK_HISTORICAL_FEATURE_OUTPUT_FORMAT) + feature_sources = [ + feature_table.batch_source for feature_table in feature_tables + ] if isinstance(entity_source, pd.DataFrame): - staging_location = self._config.get(CONFIG_SPARK_STAGING_LOCATION) - entity_staging_uri = urlparse( - os.path.join(staging_location, str(uuid.uuid4())) - ) - staging_client = get_staging_client(entity_staging_uri.scheme) - with tempfile.NamedTemporaryFile() as df_export_path: - entity_source.to_parquet(df_export_path.name) - bucket = ( - None - if entity_staging_uri.scheme == "file" - else entity_staging_uri.netloc + if any(isinstance(source, BigQuerySource) for source in feature_sources): + first_bq_source = [ + source + for source in feature_sources + if isinstance(source, BigQuerySource) + ][0] + source_ref = table_reference_from_string( + first_bq_source.bigquery_options.table_ref ) - staging_client.upload_file( - df_export_path.name, bucket, entity_staging_uri.path.lstrip("/") + entity_source = stage_entities_to_bq( + entity_source, source_ref.project, source_ref.dataset_id ) - entity_source = FileSource( - "event_timestamp", ParquetFormat(), entity_staging_uri.geturl(), + else: + entity_source = stage_entities_to_fs( + entity_source, + staging_location=self._config.get(CONFIG_SPARK_STAGING_LOCATION), ) if self._use_job_service: @@ -943,6 +951,7 @@ def get_historical_features( feature_refs=feature_refs, entity_source=entity_source.to_proto(), project=project, + output_format=output_format, output_location=output_location, ), **self._extra_grpc_params(), @@ -955,11 +964,7 @@ def get_historical_features( ) else: return start_historical_feature_retrieval_job( - self, - entity_source, - feature_tables, - output_format, - os.path.join(output_location, str(uuid.uuid4())), + self, entity_source, feature_tables, output_format, output_location, ) def get_historical_features_df( diff --git a/sdk/python/feast/job_service.py b/sdk/python/feast/job_service.py index 4fe1d22b4e5..1587dc5cb68 100644 --- a/sdk/python/feast/job_service.py +++ b/sdk/python/feast/job_service.py @@ -7,6 +7,7 @@ from feast.core import JobService_pb2_grpc from feast.core.JobService_pb2 import ( CancelJobResponse, + GetHistoricalFeaturesRequest, GetHistoricalFeaturesResponse, GetJobResponse, ) @@ -20,6 +21,7 @@ SparkJobStatus, StreamIngestionJob, ) +from feast.pyspark.launcher import start_historical_feature_retrieval_job from feast.third_party.grpc.health.v1 import HealthService_pb2_grpc from feast.third_party.grpc.health.v1.HealthService_pb2 import ( HealthCheckResponse, @@ -64,13 +66,16 @@ def StartOfflineToOnlineIngestionJob(self, request, context): context.set_details("Method not implemented!") raise NotImplementedError("Method not implemented!") - def GetHistoricalFeatures(self, request, context): + def GetHistoricalFeatures(self, request: GetHistoricalFeaturesRequest, context): """Produce a training dataset, return a job id that will provide a file reference""" - job = self.client.get_historical_features( - request.feature_refs, + job = start_historical_feature_retrieval_job( + client=self.client, entity_source=DataSource.from_proto(request.entity_source), - project=request.project, - output_location=request.output_location, + feature_tables=self.client._get_feature_tables_from_feature_refs( + list(request.feature_refs), request.project + ), + output_format=request.output_format, + output_path=request.output_location, ) output_file_uri = job.get_output_file_uri(block=False) diff --git a/sdk/python/feast/pyspark/historical_feature_retrieval_job.py b/sdk/python/feast/pyspark/historical_feature_retrieval_job.py index 67a05107b25..a27ac3406b6 100644 --- a/sdk/python/feast/pyspark/historical_feature_retrieval_job.py +++ b/sdk/python/feast/pyspark/historical_feature_retrieval_job.py @@ -149,25 +149,31 @@ def spark_format(self) -> str: def spark_path(self) -> str: return f"{self.project}:{self.dataset}.{self.table}" + @property + def spark_read_options(self) -> Dict[str, str]: + return {**super().spark_read_options, "viewsEnabled": "true"} + def _source_from_dict(dct: Dict) -> Source: if "file" in dct.keys(): return FileSource( - FileSource.PROTO_FORMAT_TO_SPARK[dct["file"]["format"]["json_class"]], - dct["file"]["path"], - dct["file"]["event_timestamp_column"], - dct["file"].get("created_timestamp_column"), - dct["file"].get("field_mapping"), - dct["file"].get("options"), + format=FileSource.PROTO_FORMAT_TO_SPARK[ + dct["file"]["format"]["json_class"] + ], + path=dct["file"]["path"], + event_timestamp_column=dct["file"]["event_timestamp_column"], + created_timestamp_column=dct["file"].get("created_timestamp_column"), + field_mapping=dct["file"].get("field_mapping"), + options=dct["file"].get("options"), ) else: return BigQuerySource( - dct["bq"]["project"], - dct["bq"]["dataset"], - dct["bq"]["table"], - dct["bq"].get("field_mapping", {}), - dct["bq"]["event_timestamp_column"], - dct["bq"].get("created_timestamp_column"), + project=dct["bq"]["project"], + dataset=dct["bq"]["dataset"], + table=dct["bq"]["table"], + field_mapping=dct["bq"].get("field_mapping", {}), + event_timestamp_column=dct["bq"]["event_timestamp_column"], + created_timestamp_column=dct["bq"].get("created_timestamp_column"), ) diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py index aeacc7b1d4e..8260383fbc4 100644 --- a/sdk/python/feast/pyspark/launcher.py +++ b/sdk/python/feast/pyspark/launcher.py @@ -35,6 +35,7 @@ StreamIngestionJob, StreamIngestionJobParameters, ) +from feast.staging.entities import create_bq_view_of_joined_features_and_entities from feast.staging.storage_client import get_staging_client from feast.value_type import ValueType @@ -106,7 +107,11 @@ def _source_to_argument(source: DataSource): return {"file": properties} if isinstance(source, BigQuerySource): - properties["table_ref"] = source.bigquery_options.table_ref + project, dataset_and_table = source.bigquery_options.table_ref.split(":") + dataset, table = dataset_and_table.split(".") + properties["project"] = project + properties["dataset"] = dataset + properties["table"] = table return {"bq": properties} if isinstance(source, KafkaSource): @@ -171,13 +176,17 @@ def start_historical_feature_retrieval_job( output_path: str, ) -> RetrievalJob: launcher = resolve_launcher(client._config) + feature_sources = [ + _source_to_argument( + replace_bq_table_with_joined_view(feature_table, entity_source) + ) + for feature_table in feature_tables + ] + return launcher.historical_feature_retrieval( RetrievalJobParameters( entity_source=_source_to_argument(entity_source), - feature_tables_sources=[ - _source_to_argument(feature_table.batch_source) - for feature_table in feature_tables - ], + feature_tables_sources=feature_sources, feature_tables=[ _feature_table_to_argument(client, feature_table) for feature_table in feature_tables @@ -188,6 +197,35 @@ def start_historical_feature_retrieval_job( ) +def replace_bq_table_with_joined_view( + feature_table: FeatureTable, entity_source: Union[FileSource, BigQuerySource], +) -> Union[FileSource, BigQuerySource]: + """ + Applies optimization to historical retrieval. Instead of pulling all data from Batch Source, + with this optimization we join feature values & entities on Data Warehouse side (improving data locality). + Several conditions should be met to enable this optimization: + * entities are staged to BigQuery + * feature values are in in BigQuery + * Entity columns are not mapped (ToDo: fix this limitation) + :return: replacement for feature source + """ + if not isinstance(feature_table.batch_source, BigQuerySource): + return feature_table.batch_source + + if not isinstance(entity_source, BigQuerySource): + return feature_table.batch_source + + if any( + entity in feature_table.batch_source.field_mapping + for entity in feature_table.entities + ): + return feature_table.batch_source + + return create_bq_view_of_joined_features_and_entities( + feature_table.batch_source, entity_source, feature_table.entities, + ) + + def _download_jar(remote_jar: str) -> str: remote_jar_parts = urlparse(remote_jar) diff --git a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py index 3b38f326e9b..66e548770ad 100644 --- a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py +++ b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py @@ -104,6 +104,8 @@ class DataprocClusterLauncher(JobLauncher): addition to the Feast SDK. """ + EXTERNAL_JARS = ["gs://spark-lib/bigquery/spark-bigquery-latest_2.12.jar"] + def __init__( self, cluster_name: str, staging_location: str, region: str, project_id: str, ): @@ -157,7 +159,7 @@ def dataproc_submit(self, job_params: SparkJobParameters) -> Operation: job_config.update( { "spark_job": { - "jar_file_uris": [main_file_uri], + "jar_file_uris": [main_file_uri] + self.EXTERNAL_JARS, "main_class": job_params.get_class_name(), "args": job_params.get_arguments(), } @@ -168,6 +170,7 @@ def dataproc_submit(self, job_params: SparkJobParameters) -> Operation: { "pyspark_job": { "main_python_file_uri": main_file_uri, + "jar_file_uris": self.EXTERNAL_JARS, "args": job_params.get_arguments(), } } diff --git a/sdk/python/feast/pyspark/launchers/standalone/local.py b/sdk/python/feast/pyspark/launchers/standalone/local.py index d7bda7a59c2..f3106c68516 100644 --- a/sdk/python/feast/pyspark/launchers/standalone/local.py +++ b/sdk/python/feast/pyspark/launchers/standalone/local.py @@ -148,6 +148,8 @@ class StandaloneClusterLauncher(JobLauncher): Submits jobs to a standalone Spark cluster in client mode. """ + BQ_CONNECTOR_VERSION = "2.12:0.17.3" + def __init__(self, master_url: str, spark_home: str = None): """ This launcher executes the spark-submit script in a subprocess. The subprocess @@ -184,6 +186,23 @@ def spark_submit( if ui_port: submission_cmd.extend(["--conf", f"spark.ui.port={ui_port}"]) + # Workaround for https://github.com/apache/spark/pull/26552 + # Fix running spark job with bigquery connector (w/ shadowing) on JDK 9+ + submission_cmd.extend( + [ + "--conf", + "spark.executor.extraJavaOptions=" + "-Dcom.google.cloud.spark.bigquery.repackaged.io.netty.tryReflectionSetAccessible=true -Duser.timezone=GMT", + "--conf", + "spark.driver.extraJavaOptions=" + "-Dcom.google.cloud.spark.bigquery.repackaged.io.netty.tryReflectionSetAccessible=true -Duser.timezone=GMT", + "--conf", + "spark.sql.session.timeZone=UTC", # ignore local timezone + "--packages", + f"com.google.cloud.spark:spark-bigquery-with-dependencies_{self.BQ_CONNECTOR_VERSION}", + ] + ) + if job_params.get_extra_options(): submission_cmd.extend(job_params.get_extra_options().split(" ")) diff --git a/sdk/python/feast/staging/entities.py b/sdk/python/feast/staging/entities.py new file mode 100644 index 00000000000..8a4745fe24c --- /dev/null +++ b/sdk/python/feast/staging/entities.py @@ -0,0 +1,129 @@ +import os +import tempfile +import uuid +from datetime import datetime, timedelta +from typing import List +from urllib.parse import urlparse + +import pandas as pd + +from feast.data_format import ParquetFormat +from feast.data_source import BigQuerySource, FileSource +from feast.staging.storage_client import get_staging_client + +try: + from google.cloud import bigquery +except ImportError: + bigquery = None + + +def stage_entities_to_fs( + entity_source: pd.DataFrame, staging_location: str +) -> FileSource: + """ + Dumps given (entities) dataframe as parquet file and stage it to remote file storage (subdirectory of staging_location) + + :return: FileSource with remote destination path + """ + entity_staging_uri = urlparse(os.path.join(staging_location, str(uuid.uuid4()))) + staging_client = get_staging_client(entity_staging_uri.scheme) + with tempfile.NamedTemporaryFile() as df_export_path: + entity_source.to_parquet(df_export_path.name) + bucket = ( + None if entity_staging_uri.scheme == "file" else entity_staging_uri.netloc + ) + staging_client.upload_file( + df_export_path.name, bucket, entity_staging_uri.path.lstrip("/") + ) + + # ToDo: support custom event_timestamp_column + return FileSource( + event_timestamp_column="event_timestamp", + file_format=ParquetFormat(), + file_url=entity_staging_uri.geturl(), + ) + + +def table_reference_from_string(table_ref: str): + """ + Parses reference string with format "{project}:{dataset}.{table}" into bigquery.TableReference + """ + project, dataset_and_table = table_ref.split(":") + dataset, table_id = dataset_and_table.split(".") + return bigquery.TableReference( + bigquery.DatasetReference(project, dataset), table_id + ) + + +def stage_entities_to_bq( + entity_source: pd.DataFrame, project: str, dataset: str +) -> BigQuerySource: + """ + Stores given (entity) dataframe as new table in BQ. Name of the table generated based on current time. + Table will expire in 1 day. + Returns BigQuerySource with reference to created table. + """ + bq_client = bigquery.Client() + destination = bigquery.TableReference( + bigquery.DatasetReference(project, dataset), + f"_entities_{datetime.now():%Y%m%d%H%M%s}", + ) + + load_job: bigquery.LoadJob = bq_client.load_table_from_dataframe( + entity_source, destination + ) + load_job.result() # wait until complete + + dest_table: bigquery.Table = bq_client.get_table(destination) + dest_table.expires = datetime.now() + timedelta(days=1) + bq_client.update_table(dest_table, fields=["expires"]) + + return BigQuerySource( + event_timestamp_column="event_timestamp", + table_ref=f"{destination.project}:{destination.dataset_id}.{destination.table_id}", + ) + + +JOIN_TEMPLATE = """SELECT + source.* +FROM + `{entities.project}.{entities.dataset_id}.{entities.table_id}` entities +JOIN + `{source.project}.{source.dataset_id}.{source.table_id}` source +ON + ({entity_key})""" + + +def create_bq_view_of_joined_features_and_entities( + source: BigQuerySource, entity_source: BigQuerySource, entity_names: List[str] +) -> BigQuerySource: + """ + Creates BQ view that joins tables from `source` and `entity_source` with join key derived from `entity_names`. + Returns BigQuerySource with reference to created view. + """ + bq_client = bigquery.Client() + + source_ref = table_reference_from_string(source.bigquery_options.table_ref) + entities_ref = table_reference_from_string(entity_source.bigquery_options.table_ref) + + destination_ref = bigquery.TableReference( + bigquery.DatasetReference(source_ref.project, source_ref.dataset_id), + f"_view_{source_ref.table_id}_{datetime.now():%Y%m%d%H%M%s}", + ) + + view = bigquery.Table(destination_ref) + view.view_query = JOIN_TEMPLATE.format( + entities=entities_ref, + source=source_ref, + entity_key=",".join([f"source.{e} = entities.{e}" for e in entity_names]), + ) + view.expires = datetime.now() + timedelta(days=1) + bq_client.create_table(view) + + return BigQuerySource( + event_timestamp_column=source.event_timestamp_column, + created_timestamp_column=source.created_timestamp_column, + table_ref=f"{view.project}:{view.dataset_id}.{view.table_id}", + field_mapping=source.field_mapping, + date_partition_column=source.date_partition_column, + ) diff --git a/sdk/python/tests/test_client.py b/sdk/python/tests/test_client.py index 54948c273b4..188e2542db0 100644 --- a/sdk/python/tests/test_client.py +++ b/sdk/python/tests/test_client.py @@ -617,7 +617,7 @@ def test_get_online_features( project="driver_project", ) # type: GetOnlineFeaturesResponse mocked_client._serving_service_stub.GetOnlineFeaturesV2.assert_called_with( - request, metadata=auth_metadata + request, metadata=auth_metadata, timeout=10 ) got_fields = got_response.field_values[1].fields @@ -707,7 +707,7 @@ def test_get_online_features_multi_entities( project="driver_project", ) # type: GetOnlineFeaturesResponse mocked_client._serving_service_stub.GetOnlineFeaturesV2.assert_called_with( - request, metadata=auth_metadata + request, metadata=auth_metadata, timeout=10 ) got_fields = got_response.field_values[1].fields diff --git a/spark/ingestion/src/main/scala/feast/ingestion/sources/bq/BigQueryReader.scala b/spark/ingestion/src/main/scala/feast/ingestion/sources/bq/BigQueryReader.scala index 05264bd55aa..1b0a57eea98 100644 --- a/spark/ingestion/src/main/scala/feast/ingestion/sources/bq/BigQueryReader.scala +++ b/spark/ingestion/src/main/scala/feast/ingestion/sources/bq/BigQueryReader.scala @@ -32,6 +32,7 @@ object BigQueryReader { ): DataFrame = { sqlContext.read .format("bigquery") + .option("viewsEnabled", "true") .load(s"${source.project}.${source.dataset}.${source.table}") .filter(col(source.eventTimestampColumn) >= new Timestamp(start.getMillis)) .filter(col(source.eventTimestampColumn) < new Timestamp(end.getMillis)) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index c96ff7c3be2..35a6c4c5658 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -22,20 +22,14 @@ def pytest_addoption(parser): parser.addoption("--redis-url", action="store", default="localhost:6379") parser.addoption("--redis-cluster", action="store_true") parser.addoption("--feast-version", action="store") - - -def pytest_runtest_makereport(item, call): - if "incremental" in item.keywords: - if call.excinfo is not None: - parent = item.parent - parent._previousfailed = item + parser.addoption("--bq-project", action="store") def pytest_runtest_setup(item): - if "incremental" in item.keywords: - previousfailed = getattr(item.parent, "_previousfailed", None) - if previousfailed is not None: - pytest.xfail("previous test failed (%s)" % previousfailed.name) + env_names = [mark.args[0] for mark in item.iter_markers(name="env")] + if env_names: + if item.config.getoption("env") not in env_names: + pytest.skip(f"test requires env in {env_names}") from .fixtures.base import project_root, project_version # noqa @@ -63,4 +57,7 @@ def pytest_runtest_setup(item): from .fixtures.external_services import ( # type: ignore # noqa feast_core, feast_serving, + enable_auth, ) + +from .fixtures.data import * # noqa diff --git a/tests/e2e/fixtures/client.py b/tests/e2e/fixtures/client.py index 366b0aa711b..2a139ad952d 100644 --- a/tests/e2e/fixtures/client.py +++ b/tests/e2e/fixtures/client.py @@ -18,6 +18,7 @@ def feast_client( feast_core: Tuple[str, int], feast_serving: Tuple[str, int], local_staging_path, + enable_auth, ): if pytestconfig.getoption("env") == "local": return Client( diff --git a/tests/e2e/fixtures/data.py b/tests/e2e/fixtures/data.py new file mode 100644 index 00000000000..287934f775e --- /dev/null +++ b/tests/e2e/fixtures/data.py @@ -0,0 +1,41 @@ +import os +import time +from datetime import datetime + +import pytest +from _pytest.fixtures import FixtureRequest +from google.cloud import bigquery + +from feast import BigQuerySource, FileSource +from feast.data_format import ParquetFormat + +__all__ = ("bq_dataset", "batch_source") + + +@pytest.fixture(scope="session") +def bq_dataset(pytestconfig): + client = bigquery.Client(project=pytestconfig.getoption("bq_project")) + timestamp = int(time.time()) + name = f"feast_e2e_{timestamp}" + client.create_dataset(name) + yield name + client.delete_dataset(name, delete_contents=True) + + +@pytest.fixture +def batch_source(local_staging_path: str, pytestconfig, request: FixtureRequest): + if pytestconfig.getoption("env") == "gcloud": + bq_project = pytestconfig.getoption("bq_project") + bq_dataset = request.getfixturevalue("bq_dataset") + return BigQuerySource( + event_timestamp_column="event_timestamp", + created_timestamp_column="created_timestamp", + table_ref=f"{bq_project}:{bq_dataset}.source_{datetime.now():%Y%m%d%H%M%s}", + ) + else: + return FileSource( + event_timestamp_column="event_timestamp", + created_timestamp_column="created_timestamp", + file_format=ParquetFormat(), + file_url=os.path.join(local_staging_path, "transactions"), + ) diff --git a/tests/e2e/fixtures/external_services.py b/tests/e2e/fixtures/external_services.py index 6929c16ec11..1350c7965bd 100644 --- a/tests/e2e/fixtures/external_services.py +++ b/tests/e2e/fixtures/external_services.py @@ -1,7 +1,7 @@ import pytest from pytest_redis.executor import NoopRedis -__all__ = ("feast_core", "feast_serving", "redis_server", "kafka_server") +__all__ = ("feast_core", "feast_serving", "redis_server", "kafka_server", "enable_auth") @pytest.fixture(scope="session") @@ -26,3 +26,8 @@ def feast_serving(pytestconfig): def kafka_server(pytestconfig): host, port = pytestconfig.getoption("kafka_brokers").split(":") return host, port + + +@pytest.fixture(scope="session") +def enable_auth(): + return False diff --git a/tests/e2e/fixtures/feast_services.py b/tests/e2e/fixtures/feast_services.py index ce7f8546917..5e9aed124c5 100644 --- a/tests/e2e/fixtures/feast_services.py +++ b/tests/e2e/fixtures/feast_services.py @@ -98,6 +98,8 @@ def feast_serving( feast_core, pytestconfig, ): + _wait_port_open(6565) # in case core is restarting with new config + jar = str( project_root / "serving" diff --git a/tests/e2e/test_historical_features.py b/tests/e2e/test_historical_features.py index a22f7b02989..6a618c92b12 100644 --- a/tests/e2e/test_historical_features.py +++ b/tests/e2e/test_historical_features.py @@ -1,5 +1,5 @@ -import os from datetime import datetime, timedelta +from typing import Union from urllib.parse import urlparse import gcsfs @@ -9,8 +9,8 @@ from pandas._testing import assert_frame_equal from pyarrow import parquet -from feast import Client, Entity, Feature, FeatureTable, FileSource, ValueType -from feast.data_format import ParquetFormat +from feast import Client, Entity, Feature, FeatureTable, ValueType +from feast.data_source import BigQuerySource, FileSource np.random.seed(0) @@ -28,33 +28,7 @@ def read_parquet(uri): raise ValueError("Unsupported scheme") -def test_historical_features(feast_client: Client, local_staging_path: str): - customer_entity = Entity( - name="user_id", description="Customer", value_type=ValueType.INT64 - ) - feast_client.apply_entity(customer_entity) - - max_age = Duration() - max_age.FromSeconds(2 * 86400) - - transactions_feature_table = FeatureTable( - name="transactions", - entities=["user_id"], - features=[ - Feature("daily_transactions", ValueType.DOUBLE), - Feature("total_transactions", ValueType.DOUBLE), - ], - batch_source=FileSource( - event_timestamp_column="event_timestamp", - created_timestamp_column="created_timestamp", - file_format=ParquetFormat(), - file_url=os.path.join(local_staging_path, "transactions"), - ), - max_age=max_age, - ) - - feast_client.apply_feature_table(transactions_feature_table) - +def generate_data(): retrieval_date = ( datetime.utcnow() .replace(hour=0, minute=0, second=0, microsecond=0) @@ -77,11 +51,6 @@ def test_historical_features(feast_client: Client, local_staging_path: str): "total_transactions": total_transactions, } ) - - feast_client.ingest(transactions_feature_table, transactions_df) - - feature_refs = ["transactions:daily_transactions"] - customer_df = pd.DataFrame( { "event_timestamp": [retrieval_date for _ in customers] @@ -89,18 +58,48 @@ def test_historical_features(feast_client: Client, local_staging_path: str): "user_id": customers + customers, } ) + return transactions_df, customer_df + + +def test_historical_features( + feast_client: Client, batch_source: Union[BigQuerySource, FileSource] +): + customer_entity = Entity( + name="user_id", description="Customer", value_type=ValueType.INT64 + ) + feast_client.apply_entity(customer_entity) + + max_age = Duration() + max_age.FromSeconds(2 * 86400) - job = feast_client.get_historical_features(feature_refs, customer_df) + transactions_feature_table = FeatureTable( + name="transactions", + entities=["user_id"], + features=[ + Feature("daily_transactions", ValueType.DOUBLE), + Feature("total_transactions", ValueType.DOUBLE), + ], + batch_source=batch_source, + max_age=max_age, + ) + + feast_client.apply_feature_table(transactions_feature_table) + + transactions_df, customers_df = generate_data() + feast_client.ingest(transactions_feature_table, transactions_df) + + feature_refs = ["transactions:daily_transactions"] + + job = feast_client.get_historical_features(feature_refs, customers_df) output_dir = job.get_output_file_uri() joined_df = read_parquet(output_dir) expected_joined_df = pd.DataFrame( { - "event_timestamp": [retrieval_date for _ in customers] - + [retrieval_outside_max_age_date for _ in customers], - "user_id": customers + customers, - "transactions__daily_transactions": daily_transactions - + [None] * len(customers), + "event_timestamp": customers_df.event_timestamp.tolist(), + "user_id": customers_df.user_id.tolist(), + "transactions__daily_transactions": transactions_df.daily_transactions.tolist() + + [None] * transactions_df.shape[0], } ) diff --git a/tests/e2e/test_online_features.py b/tests/e2e/test_online_features.py index 5533ad8f44a..c1dedd1e59d 100644 --- a/tests/e2e/test_online_features.py +++ b/tests/e2e/test_online_features.py @@ -4,16 +4,20 @@ import time import uuid from datetime import datetime, timedelta +from typing import Union import avro.schema import numpy as np import pandas as pd +import pytest import pytz from avro.io import BinaryEncoder, DatumWriter +from google.cloud import bigquery from kafka.admin import KafkaAdminClient from kafka.producer import KafkaProducer from feast import ( + BigQuerySource, Client, Entity, Feature, @@ -39,19 +43,16 @@ def generate_data(): return df -def test_offline_ingestion(feast_client: Client, local_staging_path: str): +def test_offline_ingestion( + feast_client: Client, batch_source: Union[BigQuerySource, FileSource] +): entity = Entity(name="s2id", description="S2id", value_type=ValueType.INT64,) feature_table = FeatureTable( name="drivers", entities=["s2id"], features=[Feature("unique_drivers", ValueType.INT64)], - batch_source=FileSource( - event_timestamp_column="event_timestamp", - created_timestamp_column="event_timestamp", - file_format=ParquetFormat(), - file_url=os.path.join(local_staging_path, "batch-storage"), - ), + batch_source=batch_source, ) feast_client.apply_entity(entity) @@ -60,25 +61,45 @@ def test_offline_ingestion(feast_client: Client, local_staging_path: str): original = generate_data() feast_client.ingest(feature_table, original) # write to batch (offline) storage - job = feast_client.start_offline_to_online_ingestion( - feature_table, datetime.today(), datetime.today() + timedelta(days=1) - ) + ingest_and_verify(feast_client, feature_table, original) - wait_retry_backoff(lambda: (None, job.get_status() == SparkJobStatus.COMPLETED), 60) - features = feast_client.get_online_features( - ["drivers:unique_drivers"], - entity_rows=[{"s2id": s2_id} for s2_id in original["s2id"].tolist()], - ).to_dict() +@pytest.mark.env("gcloud") +def test_offline_ingestion_from_bq_view(pytestconfig, bq_dataset, feast_client: Client): + original = generate_data() + bq_project = pytestconfig.getoption("bq_project") - ingested = pd.DataFrame.from_dict(features) - pd.testing.assert_frame_equal( - ingested[["s2id", "drivers:unique_drivers"]], - original[["s2id", "unique_drivers"]].rename( - columns={"unique_drivers": "drivers:unique_drivers"} + bq_client = bigquery.Client(project=bq_project) + source_ref = bigquery.TableReference( + bigquery.DatasetReference(bq_project, bq_dataset), + f"ingestion_source_{datetime.now():%Y%m%d%H%M%s}", + ) + bq_client.load_table_from_dataframe(original, source_ref).result() + + view_ref = bigquery.TableReference( + bigquery.DatasetReference(bq_project, bq_dataset), + f"ingestion_view_{datetime.now():%Y%m%d%H%M%s}", + ) + view = bigquery.Table(view_ref) + view.view_query = f"select * from `{source_ref.project}.{source_ref.dataset_id}.{source_ref.table_id}`" + bq_client.create_table(view) + + entity = Entity(name="s2id", description="S2id", value_type=ValueType.INT64) + feature_table = FeatureTable( + name="bq_ingestion", + entities=["s2id"], + features=[Feature("unique_drivers", ValueType.INT64)], + batch_source=BigQuerySource( + event_timestamp_column="event_timestamp", + table_ref=f"{view_ref.project}:{view_ref.dataset_id}.{view_ref.table_id}", ), ) + feast_client.apply_entity(entity) + feast_client.apply_feature_table(feature_table) + + ingest_and_verify(feast_client, feature_table, original) + def test_streaming_ingestion( feast_client: Client, local_staging_path: str, kafka_server @@ -153,6 +174,31 @@ def get_online_features(): ) +def ingest_and_verify( + feast_client: Client, feature_table: FeatureTable, original: pd.DataFrame +): + job = feast_client.start_offline_to_online_ingestion( + feature_table, + original.event_timestamp.min().to_pydatetime(), + original.event_timestamp.max().to_pydatetime() + timedelta(seconds=1), + ) + + wait_retry_backoff(lambda: (None, job.get_status() == SparkJobStatus.COMPLETED), 60) + + features = feast_client.get_online_features( + [f"{feature_table.name}:unique_drivers"], + entity_rows=[{"s2id": s2_id} for s2_id in original["s2id"].tolist()], + ).to_dict() + + ingested = pd.DataFrame.from_dict(features) + pd.testing.assert_frame_equal( + ingested[["s2id", f"{feature_table.name}:unique_drivers"]], + original[["s2id", "unique_drivers"]].rename( + columns={"unique_drivers": f"{feature_table.name}:unique_drivers"} + ), + ) + + def avro_schema(): return json.dumps( { diff --git a/tests/e2e/test_register.py b/tests/e2e/test_register.py index 0c89ee69cef..5c3cc46bcdb 100644 --- a/tests/e2e/test_register.py +++ b/tests/e2e/test_register.py @@ -1,11 +1,11 @@ -import os -import uuid from datetime import datetime import numpy as np import pandas as pd import pytest import pytz +from google.api_core.exceptions import NotFound +from google.cloud import bigquery from google.protobuf.duration_pb2 import Duration from pandas.testing import assert_frame_equal @@ -18,15 +18,6 @@ from feast.value_type import ValueType from feast.wait import wait_retry_backoff -DIR_PATH = os.path.dirname(os.path.realpath(__file__)) -PROJECT_NAME = "basic_" + uuid.uuid4().hex.upper()[0:6] -SUFFIX = str(int(datetime.now().timestamp())) - - -@pytest.fixture -def bq_table_id(): - return f"kf-feast:feaste2e.table{SUFFIX}" - @pytest.fixture def customer_entity(): @@ -89,7 +80,7 @@ def basic_featuretable(): @pytest.fixture -def bq_dataset(): +def bq_dataframe(): N_ROWS = 100 time_offset = datetime.utcnow().replace(tzinfo=pytz.utc) return pd.DataFrame( @@ -101,25 +92,6 @@ def bq_dataset(): ) -@pytest.fixture -def bq_featuretable(bq_table_id): - batch_source = BigQuerySource( - table_ref=bq_table_id, - event_timestamp_column="datetime", - created_timestamp_column="timestamp", - ) - return FeatureTable( - name="basic_featuretable", - entities=["driver_id", "customer_id"], - features=[ - Feature(name="dev_feature_float", dtype=ValueType.FLOAT), - Feature(name="dev_feature_string", dtype=ValueType.STRING), - ], - max_age=Duration(seconds=3600), - batch_source=batch_source, - ) - - @pytest.fixture def alltypes_entity(): return Entity( @@ -236,44 +208,54 @@ def test_get_list_alltypes( assert actual_list_feature_table == alltypes_featuretable -@pytest.mark.bq -def test_ingest( +@pytest.mark.env("gcloud") +def test_ingest_into_bq( feast_client: Client, customer_entity: Entity, driver_entity: Entity, - bq_featuretable: FeatureTable, - bq_dataset: pd.DataFrame, - bq_table_id: str, + bq_dataframe: pd.DataFrame, + bq_dataset: str, + pytestconfig, ): - gcp_project, _ = bq_table_id.split(":") - bq_table_id = bq_table_id.replace(":", ".") + bq_project = pytestconfig.getoption("bq_project") + bq_table_id = f"bq_staging_{datetime.now():%Y%m%d%H%M%s}" + ft = FeatureTable( + name="basic_featuretable", + entities=["driver_id", "customer_id"], + features=[ + Feature(name="dev_feature_float", dtype=ValueType.FLOAT), + Feature(name="dev_feature_string", dtype=ValueType.STRING), + ], + max_age=Duration(seconds=3600), + batch_source=BigQuerySource( + table_ref=f"{bq_project}:{bq_dataset}.{bq_table_id}", + event_timestamp_column="datetime", + created_timestamp_column="timestamp", + ), + ) # ApplyEntity feast_client.apply_entity(customer_entity) feast_client.apply_entity(driver_entity) # ApplyFeatureTable - feast_client.apply_feature_table(bq_featuretable) - feast_client.ingest(bq_featuretable, bq_dataset, timeout=120) - - from google.api_core.exceptions import NotFound - from google.cloud import bigquery + feast_client.apply_feature_table(ft) + feast_client.ingest(ft, bq_dataframe, timeout=120) - bq_client = bigquery.Client(project=gcp_project) + bq_client = bigquery.Client(project=bq_project) # Poll BQ for table until the table has been created def try_get_table(): - table_exist = False - table_resp = None try: - table_resp = bq_client.get_table(bq_table_id) - - if table_resp and table_resp.table_id == bq_table_id.split(".")[-1]: - table_exist = True + table = bq_client.get_table( + bigquery.TableReference( + bigquery.DatasetReference(bq_project, bq_dataset), bq_table_id + ) + ) except NotFound: - pass - - return table_resp, table_exist + return None, False + else: + return table, True wait_retry_backoff( retry_fn=try_get_table, @@ -281,11 +263,9 @@ def try_get_table(): timeout_msg="Timed out trying to get bigquery table", ) - query_string = f"SELECT * FROM `{bq_table_id}`" + query_string = f"SELECT * FROM `{bq_project}.{bq_dataset}.{bq_table_id}`" job = bq_client.query(query_string) query_df = job.to_dataframe() - assert_frame_equal(query_df, bq_dataset) - - bq_client.delete_table(bq_table_id, not_found_ok=True) + assert_frame_equal(query_df, bq_dataframe) diff --git a/tests/integration/test_simple.py b/tests/integration/test_simple.py new file mode 100644 index 00000000000..70ce8efe433 --- /dev/null +++ b/tests/integration/test_simple.py @@ -0,0 +1,2 @@ +def test_success(): + pass diff --git a/tests/pytest.ini b/tests/pytest.ini index 0e44395b67f..9adde5ec197 100644 --- a/tests/pytest.ini +++ b/tests/pytest.ini @@ -1,6 +1,8 @@ [pytest] +timeout=600 + filterwarnings = ignore::DeprecationWarning markers = - incremental: Skip subsequent tests if the previous test failed. + env: Skip test if its required environment is different from passed option --env