Add validations when materializing from file sources (#1615)

achals · web-flow · commit 99ee2ce656a2 · 2021-06-02T22:37:46.000-07:00
* Validate join keys when materializing from file sources

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Dedupe columns when extracting from the dataframe

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/file.py b/sdk/python/feast/infra/offline_stores/file.py
@@ -212,6 +212,13 @@ def pull_latest_from_table_or_query(
                 created_timestamp_column
             ].apply(lambda x: x if x.tzinfo is not None else x.replace(tzinfo=pytz.utc))
 
+        source_columns = set(source_df.columns)
+        if not set(join_key_columns).issubset(source_columns):
+            raise ValueError(
+                f"The DataFrame must have at least {set(join_key_columns)} columns present, "
+                f"but these were missing: {set(join_key_columns)- source_columns} "
+            )
+
         ts_columns = (
             [event_timestamp_column, created_timestamp_column]
             if created_timestamp_column
@@ -229,8 +236,7 @@ def pull_latest_from_table_or_query(
         # make driver_id a normal column again
         last_values_df.reset_index(inplace=True)
 
-        table = pyarrow.Table.from_pandas(
-            last_values_df[join_key_columns + feature_name_columns + ts_columns]
-        )
+        columns_to_extract = set(join_key_columns + feature_name_columns + ts_columns)
+        table = pyarrow.Table.from_pandas(last_values_df[columns_to_extract])
 
         return table
diff --git a/sdk/python/feast/registry.py b/sdk/python/feast/registry.py
@@ -132,7 +132,7 @@ def get_entity(self, name: str, project: str, allow_cache: bool = False) -> Enti
         for entity_proto in registry_proto.entities:
             if entity_proto.spec.name == name and entity_proto.spec.project == project:
                 return Entity.from_proto(entity_proto)
-        raise EntityNotFoundException(project, name)
+        raise EntityNotFoundException(name, project=project)
 
     def apply_feature_table(self, feature_table: FeatureTable, project: str):
         """
