add special case: rename

pyalex · pyalex · commit cfb32f29dcbd · 2021-09-27T10:52:56.000+08:00
Signed-off-by: Oleksii Moskalenko &lt;moskalenko.alexey@gmail.com&gt;
diff --git a/python/feast_spark/pyspark/historical_feature_retrieval_job.py b/python/feast_spark/pyspark/historical_feature_retrieval_job.py
@@ -322,13 +322,16 @@ class FileDestination(NamedTuple):
 
 def _map_column(df: DataFrame, col_mapping: Dict[str, str]):
     source_to_alias_map = {v: k for k, v in col_mapping.items()}
-    projection = [
-        expr(col_mapping.get(col_name, col_name)).alias(
-            source_to_alias_map.get(col_name, col_name)
-        )
-        for col_name in df.columns
-    ]
-    return df.select(projection)
+    projection = {}
+
+    for col_name in df.columns + list(set(col_mapping) - set(df.columns)):
+        if col_name in source_to_alias_map:
+            # column rename
+            projection[source_to_alias_map.get(col_name)] = col(col_name)
+        else:
+            projection[col_name] = expr(col_mapping.get(col_name, col_name))
+
+    return df.select([c.alias(a) for a, c in projection.items()])
 
 
 def as_of_join(
diff --git a/python/tests/test_historical_feature_retrieval.py b/python/tests/test_historical_feature_retrieval.py
@@ -47,7 +47,7 @@ def large_entity_csv_file(pytestconfig, spark):
     file_path = os.path.join(temp_dir, "large_entity")
     entity_schema = StructType(
         [
-            StructField("customer_id", IntegerType()),
+            StructField("id", IntegerType()),
             StructField("event_timestamp", TimestampType()),
         ]
     )

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ def large_entity_csv_file(pytestconfig, spark):`
`47`	`47`	`file_path = os.path.join(temp_dir, "large_entity")`
`48`	`48`	`entity_schema = StructType(`
`49`	`49`	`[`
`50`		`- StructField("customer_id", IntegerType()),`
	`50`	`+ StructField("id", IntegerType()),`
`51`	`51`	`StructField("event_timestamp", TimestampType()),`
`52`	`52`	`]`
`53`	`53`	`)`