feat: Enable Arrow-based columnar data transfers when to pandas in sparksource retrieval job

ElliotNguyen68 · ElliotNguyen68 · commit 001354c1f97e · 2024-03-09T11:02:18.000+07:00
Signed-off-by: tanlocnguyen &lt;tanlocnguyen296@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
@@ -338,6 +338,11 @@ def to_spark_df(self) -> pyspark.sql.DataFrame:
 
     def _to_df_internal(self, timeout: Optional[int] = None) -> pd.DataFrame:
         """Return dataset as Pandas DataFrame synchronously"""
+        spark_session = get_spark_session_or_start_new_with_repoconfig(
+            self._config.offline_store
+        )
+        spark_session.conf.set("spark.sql.execution.arrow.fallback.enabled", "true")
+        spark_session.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
         return self.to_spark_df().toPandas()
 
     def _to_arrow_internal(self, timeout: Optional[int] = None) -> pyarrow.Table:
