add wait_for_online_ingestion

logicalclocks · Dec 13, 2024 · f5b1206 · f5b1206
1 parent f7b4800
commit f5b1206
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/python/hsfs/engine/python.py b/python/hsfs/engine/python.py
@@ -1558,6 +1558,11 @@ def _write_dataframe_kafka(
                 ),
                 await_termination=offline_write_options.get("wait_for_job", False),
             )
+
+        # wait for online ingestion
+        if feature_group.online_enabled and offline_write_options.get("wait_for_online_ingestion", False):
+            feature_group.get_latest_online_ingestion().wait_for_completion()
+
         return feature_group.materialization_job
 
     @staticmethod

diff --git a/python/hsfs/engine/spark.py b/python/hsfs/engine/spark.py
@@ -518,6 +518,10 @@ def save_stream_dataframe(
         if await_termination:
             query.awaitTermination(timeout)
 
+            # wait for online ingestion
+            if feature_group.online_enabled and write_options.get("wait_for_online_ingestion", False):
+                feature_group.get_latest_online_ingestion().wait_for_completion()
+
         return query
 
     def _save_offline_dataframe(
@@ -573,6 +577,10 @@ def _save_online_dataframe(self, feature_group, dataframe, write_options):
             .save()
         )
 
+        # wait for online ingestion
+        if feature_group.online_enabled and write_options.get("wait_for_online_ingestion", False):
+            feature_group.get_latest_online_ingestion().wait_for_completion()
+
     def _get_headers(
         self,
         feature_group: Union[fg_mod.FeatureGroup, fg_mod.ExternalFeatureGroup],