diff --git a/daft/runners/progress_bar.py b/daft/runners/progress_bar.py index e1498c242a..e6b601500f 100644 --- 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 index cf7fa17e3a..d20cd32bcb 100644 --- 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, @@ -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(()) }) }