feast-dev · pyalex · Jan 8, 2021 · Dec 9, 2020 · Dec 9, 2020 · Dec 23, 2020
@@ -104,6 +104,15 @@ def get_class_name(self) -> Optional[str]:
         """
         return None
 
+    def get_extra_packages(self) -> List[str]:
+        """
+        Getter for extra maven packages to be included on driver and executor
+        classpath if applicable.
+        Returns:
+            List[str]: List of maven packages
+        """
+        return []
+
     @abc.abstractmethod
     def get_arguments(self) -> List[str]:
         """
@@ -122,6 +131,7 @@ def __init__(
         feature_tables_sources: List[Dict],
         entity_source: Dict,
         destination: Dict,
+        extra_packages: Optional[List[str]] = None,
     ):
         """
         Args:
@@ -130,6 +140,8 @@ def __init__(
             feature_tables (List[Dict]): List of feature table specification.
                 The order of the feature table must correspond to that of feature_tables_sources.
             destination (Dict): Retrieval job output destination.
+            extra_packages (Optional[List[str]): Extra maven packages to be included on Spark driver
+                and executors classpath.
 
         Examples:
             >>> # Entity source from file
@@ -233,6 +245,7 @@ def __init__(
         self._feature_tables_sources = feature_tables_sources
         self._entity_source = entity_source
         self._destination = destination
+        self._extra_packages = extra_packages if extra_packages else []
 
     def get_name(self) -> str:
         all_feature_tables_names = [ft["name"] for ft in self._feature_tables]
@@ -246,6 +259,9 @@ def get_main_file_path(self) -> str:
             os.path.dirname(__file__), "historical_feature_retrieval_job.py"
         )
 
+    def get_extra_packages(self) -> List[str]:
+        return self._extra_packages
+
     def get_arguments(self) -> List[str]:
         def json_b64_encode(obj) -> str:
             return b64encode(json.dumps(obj).encode("utf8")).decode("ascii")

@@ -7,6 +7,7 @@
 
 from pyspark.sql import DataFrame, SparkSession, Window
 from pyspark.sql.functions import col, expr, monotonically_increasing_id, row_number
+from pyspark.sql.types import LongType
 
 EVENT_TIMESTAMP_ALIAS = "event_timestamp"
 CREATED_TIMESTAMP_ALIAS = "created_timestamp"
@@ -728,7 +729,15 @@ def start_job(
     result = retrieve_historical_features(
         spark, entity_source_conf, feature_tables_sources_conf, feature_tables_conf
     )
+
     destination = FileDestination(**destination_conf)
+    if destination.format == "tfrecord":
+        entity_source = _source_from_dict(entity_source_conf)
+        result = result.withColumn(
+            entity_source.event_timestamp_column,
+            col(entity_source.event_timestamp_column).cast(LongType()),
+        )
+
     result.write.format(destination.format).mode("overwrite").save(destination.path)
 
 

@@ -191,6 +191,10 @@ def start_historical_feature_retrieval_job(
         for feature_table in feature_tables
     ]
 
+    extra_packages = []
+    if output_format == "tfrecord":
+        extra_packages.append("com.linkedin.sparktfrecord:spark-tfrecord_2.12:0.3.0")
+
     return launcher.historical_feature_retrieval(
         RetrievalJobParameters(
             entity_source=_source_to_argument(entity_source, client._config),
@@ -200,6 +204,7 @@ def start_historical_feature_retrieval_job(
                 for feature_table in feature_tables
             ],
             destination={"format": output_format, "path": output_path},
+            extra_packages=extra_packages,
         )
     )
 

@@ -231,6 +231,7 @@ def historical_feature_retrieval(
             pyspark_script_path,
             args=job_params.get_arguments(),
             output_file_uri=job_params.get_destination_path(),
+            packages=job_params.get_extra_packages(),
         )
 
         job_ref = self._submit_emr_job(step)

@@ -281,7 +281,10 @@ def _cancel_job(emr_client, job: EmrJobRef):
 
 
 def _historical_retrieval_step(
-    pyspark_script_path: str, args: List[str], output_file_uri: str,
+    pyspark_script_path: str,
+    args: List[str],
+    output_file_uri: str,
+    packages: List[str] = None,
 ) -> Dict[str, Any]:
 
     return {
@@ -297,7 +300,10 @@ def _historical_retrieval_step(
                     "Value": output_file_uri,
                 },
             ],
-            "Args": ["spark-submit", pyspark_script_path] + args,
+            "Args": ["spark-submit"]
+            + (["--packages", ",".join(packages)] if packages else [])
+            + [pyspark_script_path]
+            + args,
             "Jar": "command-runner.jar",
         },
     }

@@ -274,12 +274,20 @@ def dataproc_submit(
             "labels": {self.JOB_TYPE_LABEL_KEY: job_params.get_job_type().name.lower()},
         }
 
+        maven_package_properties = {
+            "spark.jars.packages": ",".join(job_params.get_extra_packages())
+        }
+        common_properties = {
+            "spark.executor.instances": self.executor_instances,
+            "spark.executor.cores": self.executor_cores,
+            "spark.executor.memory": self.executor_memory,
+        }
         # Add job hash to labels only for the stream ingestion job
         if isinstance(job_params, StreamIngestionJobParameters):
             job_config["labels"][self.JOB_HASH_LABEL_KEY] = job_params.get_job_hash()
 
         if job_params.get_class_name():
-            properties = {
+            scala_job_properties = {
                 "spark.yarn.user.classpath.first": "true",
                 "spark.executor.instances": self.executor_instances,
                 "spark.executor.cores": self.executor_cores,
@@ -288,15 +296,18 @@ def dataproc_submit(
                 "spark.pyspark.python": "python3.7",
             }
 
-            properties.update(extra_properties)
-
             job_config.update(
                 {
                     "spark_job": {
                         "jar_file_uris": [main_file_uri] + self.EXTERNAL_JARS,
                         "main_class": job_params.get_class_name(),
                         "args": job_params.get_arguments(),
-                        "properties": properties,
+                        "properties": {
+                            **scala_job_properties,
+                            **common_properties,
+                            **maven_package_properties,
+                            **extra_properties,
+                        },
                     }
                 }
             )
@@ -307,7 +318,11 @@ def dataproc_submit(
                         "main_python_file_uri": main_file_uri,
                         "jar_file_uris": self.EXTERNAL_JARS,
                         "args": job_params.get_arguments(),
-                        "properties": extra_properties if extra_properties else {},
+                        "properties": {
+                            **common_properties,
+                            **maven_package_properties,
+                            **extra_properties,
+                        },
                     }
                 }
             )

@@ -272,7 +272,7 @@ def spark_submit(
                 "--conf",
                 "spark.sql.session.timeZone=UTC",  # ignore local timezone
                 "--packages",
-                BQ_SPARK_PACKAGE,
+                ",".join([BQ_SPARK_PACKAGE] + job_params.get_extra_packages()),
                 "--jars",
                 "https://storage.googleapis.com/hadoop-lib/gcs/gcs-connector-hadoop2-latest.jar,"
                 "https://repo1.maven.org/maven2/org/apache/hadoop/hadoop-aws/2.7.3/hadoop-aws-2.7.3.jar,"

@@ -28,6 +28,7 @@
 from feast import Client, Entity, Feature, FeatureTable, FileSource, ValueType
 from feast.core import CoreService_pb2_grpc as Core
 from feast.data_format import ParquetFormat
+from feast.pyspark.abc import SparkJobStatus
 from tests.feast_core_server import CoreServicer
 
 
@@ -107,6 +108,26 @@ def client_with_local_spark(tmpdir):
     )
 
 
+@pytest.fixture()
+def client_with_tfrecord_output(tmpdir):
+    import pyspark
+
+    spark_staging_location = f"file://{os.path.join(tmpdir, 'staging')}"
+    historical_feature_output_location = (
+        f"file://{os.path.join(tmpdir, 'historical_feature_retrieval_tfrecord_output')}"
+    )
+
+    return Client(
+        core_url=f"localhost:{free_port}",
+        spark_launcher="standalone",
+        spark_standalone_master="local",
+        spark_home=os.path.dirname(pyspark.__file__),
+        spark_staging_location=spark_staging_location,
+        historical_feature_output_location=historical_feature_output_location,
+        historical_feature_output_format="tfrecord",
+    )
+
+
 @pytest.fixture()
 def driver_entity(client):
     return client.apply(Entity("driver_id", "description", ValueType.INT32))
@@ -466,3 +487,39 @@ def test_historical_feature_retrieval_with_pandas_dataframe_input(
             by=["customer_id", "driver_id", "event_timestamp"]
         ).reset_index(drop=True),
     )
+
+
+@pytest.mark.usefixtures(
+    "driver_entity",
+    "customer_entity",
+    "bookings_feature_table",
+    "transactions_feature_table",
+)
+def test_historical_feature_retrieval_with_tfrecord_output(
+    client_with_tfrecord_output,
+):
+
+    customer_driver_pairs_pandas_df = pd.DataFrame(
+        np.array(
+            [
+                [1001, 8001, datetime(year=2020, month=9, day=1, tzinfo=utc)],
+                [2001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)],
+                [2001, 8002, datetime(year=2020, month=9, day=1, tzinfo=utc)],
+                [1001, 8001, datetime(year=2020, month=9, day=2, tzinfo=utc)],
+                [1001, 8001, datetime(year=2020, month=9, day=3, tzinfo=utc)],
+                [1001, 8001, datetime(year=2020, month=9, day=4, tzinfo=utc)],
+            ]
+        ),
+        columns=["customer_id", "driver_id", "event_timestamp"],
+    )
+    customer_driver_pairs_pandas_df = customer_driver_pairs_pandas_df.astype(
+        {"customer_id": "int32", "driver_id": "int32"}
+    )
+
+    job_output = client_with_tfrecord_output.get_historical_features(
+        ["transactions:total_transactions", "bookings:total_completed_bookings"],
+        customer_driver_pairs_pandas_df,
+    )
+
+    job_output.get_output_file_uri()
+    assert job_output.get_status() == SparkJobStatus.COMPLETED
diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py
@@ -43,6 +43,7 @@ def pytest_runtest_setup(item):
     global_staging_path,
     ingestion_job_jar,
     local_staging_path,
+    tfrecord_feast_client,
 )
 
 if not os.environ.get("DISABLE_SERVICE_FIXTURES"):

diff --git a/tests/e2e/fixtures/client.py b/tests/e2e/fixtures/client.py
@@ -103,6 +103,82 @@ def feast_client(
     return c
 
 
+@pytest.fixture
+def tfrecord_feast_client(
+    pytestconfig,
+    feast_core: Tuple[str, int],
+    local_staging_path,
+    feast_jobservice: Optional[Tuple[str, int]],
+    enable_auth,
+):
+    if feast_jobservice is None:
+        job_service_env = dict()
+    else:
+        job_service_env = dict(
+            job_service_url=f"{feast_jobservice[0]}:{feast_jobservice[1]}"
+        )
+
+    if pytestconfig.getoption("env") == "local":
+        import pyspark
+
+        return Client(
+            core_url=f"{feast_core[0]}:{feast_core[1]}",
+            spark_launcher="standalone",
+            spark_standalone_master="local",
+            spark_home=os.getenv("SPARK_HOME") or os.path.dirname(pyspark.__file__),
+            spark_staging_location=os.path.join(local_staging_path, "spark"),
+            historical_feature_output_format="tfrecord",
+            historical_feature_output_location=os.path.join(
+                local_staging_path, "historical_output"
+            ),
+            **job_service_env,
+        )
+
+    elif pytestconfig.getoption("env") == "gcloud":
+        c = Client(
+            core_url=f"{feast_core[0]}:{feast_core[1]}",
+            spark_launcher="dataproc",
+            dataproc_cluster_name=pytestconfig.getoption("dataproc_cluster_name"),
+            dataproc_project=pytestconfig.getoption("dataproc_project"),
+            dataproc_region=pytestconfig.getoption("dataproc_region"),
+            spark_staging_location=os.path.join(local_staging_path, "dataproc"),
+            historical_feature_output_format="tfrecord",
+            historical_feature_output_location=os.path.join(
+                local_staging_path, "historical_output"
+            ),
+            ingestion_drop_invalid_rows=True,
+            **job_service_env,
+        )
+    elif pytestconfig.getoption("env") == "aws":
+        return Client(
+            core_url=f"{feast_core[0]}:{feast_core[1]}",
+            spark_launcher="emr",
+            emr_cluster_id=pytestconfig.getoption("emr_cluster_id"),
+            emr_region=pytestconfig.getoption("emr_region"),
+            spark_staging_location=os.path.join(local_staging_path, "emr"),
+            emr_log_location=os.path.join(local_staging_path, "emr_logs"),
+            historical_feature_output_format="tfrecord",
+            historical_feature_output_location=os.path.join(
+                local_staging_path, "historical_output"
+            ),
+        )
+    elif pytestconfig.getoption("env") == "k8s":
+        return Client(
+            core_url=f"{feast_core[0]}:{feast_core[1]}",
+            spark_launcher="k8s",
+            spark_staging_location=os.path.join(local_staging_path, "k8s"),
+            historical_feature_output_format="tfrecord",
+            historical_feature_output_location=os.path.join(
+                local_staging_path, "historical_output"
+            ),
+        )
+    else:
+        raise KeyError(f"Unknown environment {pytestconfig.getoption('env')}")
+
+    c.set_project(pytestconfig.getoption("feast_project"))
+    return c
+
+
 @pytest.fixture(scope="session")
 def global_staging_path(pytestconfig):
     if pytestconfig.getoption("env") == "local" and not pytestconfig.getoption(

diff --git a/tests/e2e/test_historical_features.py b/tests/e2e/test_historical_features.py
@@ -11,6 +11,7 @@
 
 from feast import Client, Entity, Feature, FeatureTable, ValueType
 from feast.data_source import BigQuerySource, FileSource
+from feast.pyspark.abc import SparkJobStatus
 
 np.random.seed(0)
 
@@ -68,7 +69,9 @@ def generate_data():
 
 
 def test_historical_features(
-    feast_client: Client, batch_source: Union[BigQuerySource, FileSource]
+    feast_client: Client,
+    tfrecord_feast_client: Client,
+    batch_source: Union[BigQuerySource, FileSource],
 ):
     customer_entity = Entity(
         name="user_id", description="Customer", value_type=ValueType.INT64
@@ -115,3 +118,7 @@ def test_historical_features(
             drop=True
         ),
     )
+
+    job = tfrecord_feast_client.get_historical_features(feature_refs, customers_df)
+    job.get_output_file_uri()
+    assert job.get_status() == SparkJobStatus.COMPLETED