feat: support hdfs staging for Spark offline store

chimeyrock999 · chimeyrock999 · commit 583275f30152 · 2025-09-26T10:52:30.000+07:00
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
@@ -442,6 +442,7 @@ def persist(
                 self.to_spark_df().write.format(file_format).saveAsTable(table_name)
         else:
             self.to_spark_df().createOrReplaceTempView(table_name)
+            
 
     def _has_remote_warehouse_in_config(self) -> bool:
         """
@@ -497,15 +498,35 @@ def to_remote_storage(self) -> List[str]:
                 return aws_utils.list_s3_files(
                     self._config.offline_store.region, output_uri
                 )
-
+            elif self._config.offline_store.staging_location.startswith("hdfs://"):
+                output_uri = os.path.join(
+                    self._config.offline_store.staging_location, str(uuid.uuid4())
+                )
+                sdf.write.parquet(output_uri)
+                spark_session = get_spark_session_or_start_new_with_repoconfig(
+                    store_config=self._config.offline_store
+                )
+                return self._list_hdfs_files(spark_session, output_uri)
             else:
                 raise NotImplementedError(
-                    "to_remote_storage is only implemented for file:// and s3:// uri schemes"
+                    "to_remote_storage is only implemented for file://, s3:// and hdfs:// uri schemes"
                 )
 
         else:
             raise NotImplementedError()
 
+    def _list_hdfs_files(self, spark_session: SparkSession, uri: str) -> List[str]:
+        jvm = spark_session._jvm
+        conf = spark_session._jsc.hadoopConfiguration()
+        path = jvm.org.apache.hadoop.fs.Path(uri)
+        fs = jvm.org.apache.hadoop.fs.FileSystem.get(path.toUri(), conf)
+        statuses = fs.listStatus(path)
+        files = []
+        for f in statuses:
+            if f.isFile():
+                files.append(f.getPath().toString())
+        return files
+
     @property
     def metadata(self) -> Optional[RetrievalMetadata]:
         """
