Implicit type conversion for entity and feature table source (#43)

khorshuheng · web-flow · commit 49142c2e8280 · 2021-03-15T00:24:51.000-07:00
* Implicit type conversion for entity and feature table source

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

* Restrict type casting only for double to float conversion

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

* Make allowable type casting configurable

Signed-off-by: Khor Shu Heng &lt;khor.heng@go-jek.com&gt;
diff --git a/python/feast_spark/pyspark/historical_feature_retrieval_job.py b/python/feast_spark/pyspark/historical_feature_retrieval_job.py
@@ -560,6 +560,22 @@ def _read_and_verify_entity_df_from_source(
     return mapped_entity_df
 
 
+def _type_casting_allowed(feature_type: str, source_col_type):
+    allowed_casting_for_source_col = {"double": ["float"]}
+
+    if feature_type == source_col_type:
+        return True
+
+    allowed_feature_type_for_casting = allowed_casting_for_source_col.get(
+        source_col_type
+    )
+
+    return (
+        allowed_feature_type_for_casting is not None
+        and feature_type in allowed_feature_type_for_casting
+    )
+
+
 def _read_and_verify_feature_table_df_from_source(
     spark: SparkSession, feature_table: FeatureTable, source: Source,
 ) -> DataFrame:
@@ -591,10 +607,16 @@ def _read_and_verify_feature_table_df_from_source(
     feature_table_dtypes = dict(mapped_source_df.dtypes)
     for field in feature_table.entities + feature_table.features:
         column_type = feature_table_dtypes.get(field.name)
+
         if column_type != field.spark_type:
-            raise SchemaError(
-                f"{field.name} should be of {field.spark_type} type, but is {column_type} instead"
-            )
+            if _type_casting_allowed(field.spark_type, column_type):
+                mapped_source_df = mapped_source_df.withColumn(
+                    field.name, col(field.name).cast(field.spark_type)
+                )
+            else:
+                raise SchemaError(
+                    f"{field.name} should be of {field.spark_type} type, but is {column_type} instead"
+                )
 
     for timestamp_column in [
         source.event_timestamp_column,
diff --git a/python/tests/data/bookings.csv b/python/tests/data/bookings.csv
@@ -0,0 +1,6 @@
+driver_id,event_timestamp,created_timestamp,completed_bookings
+8001,2020-08-31T00:00:00.000,2020-08-31T00:00:00.000,200
+8001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,300
+8002,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,600
+8002,2020-09-01T00:00:00.000,2020-09-02T00:00:00.000,500
+8003,2020-09-01T00:00:00.000,2020-09-02T00:00:00.000,700
diff --git a/python/tests/data/column_mapping_test_entity.csv b/python/tests/data/column_mapping_test_entity.csv
@@ -0,0 +1,6 @@
+id,event_timestamp
+1001,2020-09-02T00:00:00.000
+1001,2020-09-03T00:00:00.000
+2001,2020-09-04T00:00:00.000
+2001,2020-09-04T00:00:00.000
+3001,2020-09-04T00:00:00.000
diff --git a/python/tests/data/column_mapping_test_feature.csv b/python/tests/data/column_mapping_test_feature.csv
@@ -0,0 +1,6 @@
+customer_id,total_bookings,datetime,created_datetime
+1001,200,2020-09-02T00:00:00.000,2020-09-02T00:00:00.000
+1001,400,2020-09-04T00:00:00.000,2020-09-02T00:00:00.000
+2001,500,2020-09-03T00:00:00.000,2020-09-01T00:00:00.000
+2001,600,2020-09-03T00:00:00.000,2020-09-02T00:00:00.000
+3001,700,2020-09-03T00:00:00.000,2020-09-03T00:00:00.000
diff --git a/python/tests/data/customer_driver_pairs.csv b/python/tests/data/customer_driver_pairs.csv
@@ -0,0 +1,6 @@
+customer_id,driver_id,event_timestamp
+1001,8001,2020-09-02T00:00:00.000
+1001,8002,2020-09-02T00:00:00.000
+1001,8002,2020-09-03T00:00:00.000
+2001,8002,2020-09-03T00:00:00.000
+2001,8002,2020-09-04T00:00:00.000
diff --git a/python/tests/data/customers.csv b/python/tests/data/customers.csv
@@ -0,0 +1,6 @@
+customer_id,event_timestamp
+1001,2020-09-02T00:00:00.000
+1002,2020-09-02T00:00:00.000
+1003,2020-09-03T00:00:00.000
+1004,2020-09-03T00:00:00.000
+1005,2020-09-04T00:00:00.000
diff --git a/python/tests/data/single_customer.csv b/python/tests/data/single_customer.csv
@@ -0,0 +1,2 @@
+customer_id,event_timestamp
+1001,2020-09-02T00:00:00.000
diff --git a/python/tests/data/transactions.csv b/python/tests/data/transactions.csv
@@ -0,0 +1,6 @@
+customer_id,event_timestamp,created_timestamp,daily_transactions
+1001,2020-08-31T00:00:00.000,2020-09-01T00:00:00.000,50.0
+1001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,100.0
+2001,2020-09-01T00:00:00.000,2020-08-31T00:00:00.000,80.0
+2001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,200.0
+3001,2020-09-01T00:00:00.000,2020-09-01T00:00:00.000,300.0
diff --git a/python/tests/test_historical_feature_retrieval.py b/python/tests/test_historical_feature_retrieval.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+customer_id,event_timestamp`
	`2`	`+1001,2020-09-02T00:00:00.000`