only clear at end, mark done with message

Eventual-Inc · Dec 19, 2024 · caedb22 · caedb22
1 parent 653bc11
commit caedb22
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/daft/runners/progress_bar.py b/daft/runners/progress_bar.py
@@ -139,10 +139,9 @@ def update_bar(self, pbar_id: int, message: str) -> None:
             del self.bar_configs[pbar_id]
         self.pbars[pbar_id].set_description_str(message)
 
-    def close_bar(self, pbar_id: int) -> None:
+    def mark_finished(self, pbar_id: int, message: str) -> None:
         if pbar_id in self.pbars:
-            self.pbars[pbar_id].close()
-            del self.pbars[pbar_id]
+            self.pbars[pbar_id].set_description_str(message)
 
     def close(self) -> None:
         for p in self.pbars.values():

diff --git a/src/daft-local-execution/src/progress_bar.rs b/src/daft-local-execution/src/progress_bar.rs
@@ -13,7 +13,7 @@ use crate::runtime_stats::RuntimeStatsContext;
 
 pub trait ProgressBar: Send + Sync {
     fn set_message(&self, message: String) -> DaftResult<()>;
-    fn close(&self) -> DaftResult<()>;
+    fn mark_finished(&self, message: String) -> DaftResult<()>;
 }
 
 pub trait ProgressBarManager {
@@ -52,8 +52,8 @@ pub struct OperatorProgressBar {
 }
 
 impl OperatorProgressBar {
-    // 100ms = 100_000_000ns
-    const UPDATE_INTERVAL: u64 = 100_000_000;
+    // 500ms = 500_000_000ns
+    const UPDATE_INTERVAL: u64 = 500_000_000;
 
     pub fn new(
         progress_bar: Box<dyn ProgressBar>,
@@ -111,7 +111,9 @@ impl OperatorProgressBar {
 
 impl Drop for OperatorProgressBar {
     fn drop(&mut self) {
-        let _ = self.inner_progress_bar.close();
+        let _ = self
+            .inner_progress_bar
+            .mark_finished("Finished".to_string());
     }
 }
 
@@ -123,8 +125,8 @@ impl ProgressBar for IndicatifProgressBar {
         Ok(())
     }
 
-    fn close(&self) -> DaftResult<()> {
-        self.0.finish_and_clear();
+    fn mark_finished(&self, message: String) -> DaftResult<()> {
+        self.0.finish_with_message(message);
         Ok(())
     }
 }
@@ -219,8 +221,8 @@ mod python {
             self.manager.update_bar(self.pb_id, message.as_str())
         }
 
-        fn close(&self) -> DaftResult<()> {
-            self.manager.close_bar(self.pb_id)
+        fn mark_finished(&self, message: String) -> DaftResult<()> {
+            self.manager.mark_finished(self.pb_id, message.as_str())
         }
     }
 
@@ -250,9 +252,10 @@ mod python {
             })
         }
 
-        fn close_bar(&self, pb_id: usize) -> DaftResult<()> {
+        fn mark_finished(&self, pb_id: usize, message: &str) -> DaftResult<()> {
             Python::with_gil(|py| {
-                self.inner.call_method1(py, "close_bar", (pb_id,))?;
+                self.inner
+                    .call_method1(py, "mark_finished", (pb_id, message))?;
                 DaftResult::Ok(())
             })
         }