fix: Fix materialization when running on Spark cluster. (#3166)

ckarwicki · web-flow · commit 175fd256e0d2 · 2022-09-08T21:21:22.000-07:00
* Fix materialization when running on Spark cluster.

When running materialization and have Spark offline store configured to use cluster (`spark.master` pointing to actual Spark master node) `self.to_spark_df().write.parquet(temp_dir, mode="overwrite")` will create parquet file in worker node but `return pq.read_table(temp_dir)` is executed on driver node and it can't read from worker. Proposed fix makes materialization work when run on Spark cluster.

Signed-off-by: ckarwicki &lt;104110169+ckarwicki-deloitte@users.noreply.github.com&gt;
Signed-off-by: ckarwicki &lt;71740096+ckarwicki@users.noreply.github.com&gt;

* Fix linter.

Signed-off-by: ckarwicki &lt;jdeveloper98@gmail.com&gt;
Signed-off-by: ckarwicki &lt;104110169+ckarwicki-deloitte@users.noreply.github.com&gt;
Signed-off-by: ckarwicki &lt;71740096+ckarwicki@users.noreply.github.com&gt;

* Fix linter.

Signed-off-by: ckarwicki &lt;jdeveloper98@gmail.com&gt;
Signed-off-by: ckarwicki &lt;104110169+ckarwicki-deloitte@users.noreply.github.com&gt;
Signed-off-by: ckarwicki &lt;71740096+ckarwicki@users.noreply.github.com&gt;

Signed-off-by: ckarwicki &lt;104110169+ckarwicki-deloitte@users.noreply.github.com&gt;
Signed-off-by: ckarwicki &lt;71740096+ckarwicki@users.noreply.github.com&gt;
Signed-off-by: ckarwicki &lt;jdeveloper98@gmail.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py
@@ -342,11 +342,7 @@ def _to_df_internal(self) -> pd.DataFrame:
 
     def _to_arrow_internal(self) -> pyarrow.Table:
         """Return dataset as pyarrow Table synchronously"""
-
-        # write to temp parquet and then load it as pyarrow table from disk
-        with tempfile.TemporaryDirectory() as temp_dir:
-            self.to_spark_df().write.parquet(temp_dir, mode="overwrite")
-            return pq.read_table(temp_dir)
+        return pyarrow.Table.from_pandas(self._to_df_internal())
 
     def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         """
