Make created timestamp column optional (#1108)

khorshuheng · web-flow · commit 4699d58bc1c1 · 2020-10-29T15:12:55.000+08:00
* Make created timestamp column optional

Signed-off-by: Khor Shu Heng &lt;khor.heng@gojek.com&gt;

* Fix argument for FileSource under e2e tests

Signed-off-by: Khor Shu Heng &lt;khor.heng@gojek.com&gt;

Co-authored-by: Khor Shu Heng &lt;khor.heng@gojek.com&gt;
diff --git a/docs/user-guide/feature-retrieval.md b/docs/user-guide/feature-retrieval.md
@@ -36,8 +36,7 @@ feature_refs = [
 # Define entity source
 entity_source = FileSource(
    "event_timestamp",
-   "created_timestamp",
-   "parquet",
+   ParquetFormat(),
    "gs://some-bucket/customer"
 )
 
diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
@@ -485,9 +485,7 @@ def get_historical_features(features: str, entity_df_path: str, destination: str
 
     entity_df["event_timestamp"] = pandas.to_datetime(entity_df["event_timestamp"])
 
-    uploaded_df = client.stage_dataframe(
-        entity_df, "event_timestamp", "created_timestamp"
-    )
+    uploaded_df = client.stage_dataframe(entity_df, "event_timestamp")
 
     job = client.get_historical_features(features.split(","), uploaded_df,)
     print(job.get_output_file_uri())
diff --git a/sdk/python/feast/client.py b/sdk/python/feast/client.py
@@ -881,10 +881,11 @@ def get_historical_features(
 
         Examples:
             >>> from feast import Client
+            >>> from feast.data_format import ParquetFormat
             >>> from datetime import datetime
             >>> feast_client = Client(core_url="localhost:6565")
             >>> feature_refs = ["bookings:bookings_7d", "bookings:booking_14d"]
-            >>> entity_source = FileSource("event_timestamp", "parquet", "gs://some-bucket/customer")
+            >>> entity_source = FileSource("event_timestamp", ParquetFormat(), "gs://some-bucket/customer")
             >>> feature_retrieval_job = feast_client.get_historical_features(
             >>>     feature_refs, entity_source, project="my_project")
             >>> output_file_uri = feature_retrieval_job.get_output_file_uri()
@@ -916,10 +917,7 @@ def get_historical_features(
                     df_export_path.name, bucket, entity_staging_uri.path.lstrip("/")
                 )
                 entity_source = FileSource(
-                    "event_timestamp",
-                    "created_timestamp",
-                    ParquetFormat(),
-                    entity_staging_uri.geturl(),
+                    "event_timestamp", ParquetFormat(), entity_staging_uri.geturl(),
                 )
 
         return start_historical_feature_retrieval_job(
@@ -955,12 +953,13 @@ def get_historical_features_df(
 
         Examples:
             >>> from feast import Client
+            >>> from feast.data_format import ParquetFormat
             >>> from datetime import datetime
             >>> from pyspark.sql import SparkSession
             >>> spark = SparkSession.builder.getOrCreate()
             >>> feast_client = Client(core_url="localhost:6565")
             >>> feature_refs = ["bookings:bookings_7d", "bookings:booking_14d"]
-            >>> entity_source = FileSource("event_timestamp", "parquet", "gs://some-bucket/customer")
+            >>> entity_source = FileSource("event_timestamp", ParquetFormat, "gs://some-bucket/customer")
             >>> df = feast_client.get_historical_features(
             >>>     feature_refs, entity_source, project="my_project")
         """
@@ -1011,11 +1010,6 @@ def start_stream_to_online_ingestion(
         return start_stream_to_online_ingestion(feature_table, extra_jars or [], self)
 
     def stage_dataframe(
-        self,
-        df: pd.DataFrame,
-        event_timestamp_column: str,
-        created_timestamp_column: str,
+        self, df: pd.DataFrame, event_timestamp_column: str,
     ) -> FileSource:
-        return stage_dataframe(
-            df, event_timestamp_column, created_timestamp_column, self
-        )
+        return stage_dataframe(df, event_timestamp_column, self)
diff --git a/sdk/python/feast/data_source.py b/sdk/python/feast/data_source.py
@@ -348,13 +348,13 @@ class DataSource:
     def __init__(
         self,
         event_timestamp_column: str,
-        created_timestamp_column: str,
-        field_mapping: Optional[Dict[str, str]] = dict(),
+        created_timestamp_column: Optional[str] = "",
+        field_mapping: Optional[Dict[str, str]] = None,
         date_partition_column: Optional[str] = "",
     ):
         self._event_timestamp_column = event_timestamp_column
         self._created_timestamp_column = created_timestamp_column
-        self._field_mapping = field_mapping
+        self._field_mapping = field_mapping if field_mapping else {}
         self._date_partition_column = date_partition_column
 
     def __eq__(self, other):
@@ -409,7 +409,7 @@ def created_timestamp_column(self):
     @created_timestamp_column.setter
     def created_timestamp_column(self, created_timestamp_column):
         """
-        Sets the event timestamp column of this data source
+        Sets the created timestamp column of this data source
         """
         self._created_timestamp_column = created_timestamp_column
 
@@ -498,10 +498,10 @@ class FileSource(DataSource):
     def __init__(
         self,
         event_timestamp_column: str,
-        created_timestamp_column: str,
         file_format: FileFormat,
         file_url: str,
-        field_mapping: Optional[Dict[str, str]] = dict(),
+        created_timestamp_column: Optional[str] = "",
+        field_mapping: Optional[Dict[str, str]] = None,
         date_partition_column: Optional[str] = "",
     ):
         super().__init__(
@@ -555,9 +555,9 @@ class BigQuerySource(DataSource):
     def __init__(
         self,
         event_timestamp_column: str,
-        created_timestamp_column: str,
         table_ref: str,
-        field_mapping: Optional[Dict[str, str]] = dict(),
+        created_timestamp_column: Optional[str] = "",
+        field_mapping: Optional[Dict[str, str]] = None,
         date_partition_column: Optional[str] = "",
     ):
         super().__init__(
diff --git a/sdk/python/feast/pyspark/abc.py b/sdk/python/feast/pyspark/abc.py
@@ -458,10 +458,7 @@ def start_stream_to_online_ingestion(
 
     @abc.abstractmethod
     def stage_dataframe(
-        self,
-        df: pandas.DataFrame,
-        event_timestamp_column: str,
-        created_timestamp_column: str,
+        self, df: pandas.DataFrame, event_timestamp_column: str,
     ) -> FileSource:
         """
         Upload a pandas dataframe so it is available to the Spark cluster.
diff --git a/sdk/python/feast/pyspark/historical_feature_retrieval_job.py b/sdk/python/feast/pyspark/historical_feature_retrieval_job.py
@@ -86,7 +86,7 @@ def __init__(
         format: str,
         path: str,
         event_timestamp_column: str,
-        created_timestamp_column: Optional[str] = None,
+        created_timestamp_column: Optional[str] = "",
         field_mapping: Optional[Dict[str, str]] = None,
         options: Optional[Dict[str, str]] = None,
     ):
diff --git a/sdk/python/feast/pyspark/launcher.py b/sdk/python/feast/pyspark/launcher.py
@@ -243,10 +243,6 @@ def start_stream_to_online_ingestion(
     )
 
 
-def stage_dataframe(
-    df, event_timestamp_column: str, created_timestamp_column: str, client: "Client"
-) -> FileSource:
+def stage_dataframe(df, event_timestamp_column: str, client: "Client") -> FileSource:
     launcher = resolve_launcher(client._config)
-    return launcher.stage_dataframe(
-        df, event_timestamp_column, created_timestamp_column,
-    )
+    return launcher.stage_dataframe(df, event_timestamp_column)
diff --git a/sdk/python/feast/pyspark/launchers/aws/emr.py b/sdk/python/feast/pyspark/launchers/aws/emr.py
@@ -291,9 +291,7 @@ def start_stream_to_online_ingestion(
 
         return EmrStreamIngestionJob(self._emr_client(), job_ref)
 
-    def stage_dataframe(
-        self, df: pandas.DataFrame, event_timestamp: str, created_timestamp_column: str
-    ) -> FileSource:
+    def stage_dataframe(self, df: pandas.DataFrame, event_timestamp: str) -> FileSource:
         with tempfile.NamedTemporaryFile() as f:
             df.to_parquet(f)
             file_url = _s3_upload(
@@ -304,7 +302,6 @@ def stage_dataframe(
             )
         return FileSource(
             event_timestamp_column=event_timestamp,
-            created_timestamp_column=created_timestamp_column,
             file_format=ParquetFormat(),
             file_url=file_url,
         )
diff --git a/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py b/sdk/python/feast/pyspark/launchers/gcloud/dataproc.py
@@ -205,9 +205,7 @@ def start_stream_to_online_ingestion(
         cancel_fn = partial(self.dataproc_cancel, operation.metadata.job_id)
         return DataprocStreamingIngestionJob(operation, cancel_fn)
 
-    def stage_dataframe(
-        self, df, event_timestamp_column: str, created_timestamp_column: str,
-    ):
+    def stage_dataframe(self, df, event_timestamp_column: str):
         raise NotImplementedError
 
     def get_job_by_id(self, job_id: str) -> SparkJob:
diff --git a/sdk/python/feast/pyspark/launchers/standalone/local.py b/sdk/python/feast/pyspark/launchers/standalone/local.py
@@ -224,9 +224,7 @@ def start_stream_to_online_ingestion(
             ui_port,
         )
 
-    def stage_dataframe(
-        self, df, event_timestamp_column: str, created_timestamp_column: str,
-    ):
+    def stage_dataframe(self, df, event_timestamp_column: str):
         raise NotImplementedError
 
     def get_job_by_id(self, job_id: str) -> SparkJob:
diff --git a/sdk/python/tests/test_historical_feature_retrieval.py b/sdk/python/tests/test_historical_feature_retrieval.py
@@ -179,7 +179,10 @@ def transactions_feature_table(spark, client):
         spark, "transactions", schema, df_data
     )
     file_source = FileSource(
-        "event_timestamp", "created_timestamp", ParquetFormat(), file_uri
+        event_timestamp_column="event_timestamp",
+        created_timestamp_column="created_timestamp",
+        file_format=ParquetFormat(),
+        file_url=file_uri,
     )
     features = [
         Feature("total_transactions", ValueType.DOUBLE),
@@ -225,7 +228,10 @@ def bookings_feature_table(spark, client):
     temp_dir, file_uri = create_temp_parquet_file(spark, "bookings", schema, df_data)
 
     file_source = FileSource(
-        "event_timestamp", "created_timestamp", ParquetFormat(), file_uri
+        event_timestamp_column="event_timestamp",
+        created_timestamp_column="created_timestamp",
+        file_format=ParquetFormat(),
+        file_url=file_uri,
     )
     features = [Feature("total_completed_bookings", ValueType.INT32)]
     max_age = Duration()
@@ -270,7 +276,11 @@ def bookings_feature_table_with_mapping(spark, client):
     temp_dir, file_uri = create_temp_parquet_file(spark, "bookings", schema, df_data)
 
     file_source = FileSource(
-        "datetime", "created_datetime", ParquetFormat(), file_uri, {"id": "driver_id"}
+        event_timestamp_column="datetime",
+        created_timestamp_column="created_datetime",
+        file_format=ParquetFormat(),
+        file_url=file_uri,
+        field_mapping={"id": "driver_id"},
     )
     features = [Feature("total_completed_bookings", ValueType.INT32)]
     max_age = Duration()
@@ -309,7 +319,9 @@ def test_historical_feature_retrieval_from_local_spark_session(
         spark, "customer_driver_pair", schema, df_data
     )
     customer_driver_pairs_source = FileSource(
-        "event_timestamp", "created_timestamp", ParquetFormat(), file_uri
+        event_timestamp_column="event_timestamp",
+        file_format=ParquetFormat(),
+        file_url=file_uri,
     )
     joined_df = client.get_historical_features_df(
         ["transactions:total_transactions", "bookings:total_completed_bookings"],
@@ -356,7 +368,9 @@ def test_historical_feature_retrieval_with_field_mappings_from_local_spark_sessi
     ]
     temp_dir, file_uri = create_temp_parquet_file(spark, "drivers", schema, df_data)
     entity_source = FileSource(
-        "event_timestamp", "created_timestamp", ParquetFormat(), file_uri
+        event_timestamp_column="event_timestamp",
+        file_format=ParquetFormat(),
+        file_url=file_uri,
     )
     joined_df = client.get_historical_features_df(
         ["bookings:total_completed_bookings"], entity_source,
diff --git a/tests/e2e/test_historical_features.py b/tests/e2e/test_historical_features.py
@@ -45,10 +45,10 @@ def test_historical_features(feast_client: Client, local_staging_path: str):
             Feature("total_transactions", ValueType.DOUBLE),
         ],
         batch_source=FileSource(
-            "event_timestamp",
-            "created_timestamp",
-            ParquetFormat(),
-            os.path.join(local_staging_path, "transactions"),
+            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,
     )
diff --git a/tests/e2e/test_online_features.py b/tests/e2e/test_online_features.py
@@ -47,10 +47,10 @@ def test_offline_ingestion(feast_client: Client, local_staging_path: str):
         entities=["s2id"],
         features=[Feature("unique_drivers", ValueType.INT64)],
         batch_source=FileSource(
-            "event_timestamp",
-            "event_timestamp",
-            ParquetFormat(),
-            os.path.join(local_staging_path, "batch-storage"),
+            event_timestamp_column="event_timestamp",
+            created_timestamp_column="event_timestamp",
+            file_format=ParquetFormat(),
+            file_url=os.path.join(local_staging_path, "batch-storage"),
         ),
     )
 
@@ -92,10 +92,10 @@ def test_streaming_ingestion(
         entities=["s2id"],
         features=[Feature("unique_drivers", ValueType.INT64)],
         batch_source=FileSource(
-            "event_timestamp",
-            "event_timestamp",
-            ParquetFormat(),
-            os.path.join(local_staging_path, "batch-storage"),
+            event_timestamp_column="event_timestamp",
+            created_timestamp_column="event_timestamp",
+            file_format=ParquetFormat(),
+            file_url=os.path.join(local_staging_path, "batch-storage"),
         ),
         stream_source=KafkaSource(
             "event_timestamp",

Original file line number	Diff line number	Diff line change
`@@ -36,8 +36,7 @@ feature_refs = [`
`36`	`36`	`# Define entity source`
`37`	`37`	`entity_source = FileSource(`
`38`	`38`	`"event_timestamp",`
`39`		`- "created_timestamp",`
`40`		`- "parquet",`
	`39`	`+ ParquetFormat(),`
`41`	`40`	`"gs://some-bucket/customer"`
`42`	`41`	`)`
`43`	`42`
Original file line number	Diff line number	Diff line change
`@@ -224,9 +224,7 @@ def start_stream_to_online_ingestion(`
`224`	`224`	`ui_port,`
`225`	`225`	`)`
`226`	`226`
`227`		`- def stage_dataframe(`
`228`		`- self, df, event_timestamp_column: str, created_timestamp_column: str,`
`229`		`- ):`
	`227`	`+ def stage_dataframe(self, df, event_timestamp_column: str):`
`230`	`228`	`raise NotImplementedError`
`231`	`229`
`232`	`230`	`def get_job_by_id(self, job_id: str) -> SparkJob:`