You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by mi...@apache.org on 2026/05/11 08:24:14 UTC
(datafusion-ballista) branch main updated: fix: rest api calculates stage running time correctly (#1675)
This is an automated email from the ASF dual-hosted git repository.
milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new 146263592 fix: rest api calculates stage running time correctly (#1675)
146263592 is described below
commit 146263592d70175385ef035fa7275544f99b95bb
Author: Marko Milenković <mi...@users.noreply.github.com>
AuthorDate: Mon May 11 09:24:09 2026 +0100
fix: rest api calculates stage running time correctly (#1675)
* fix: rest api calculates stage running time correctly
* move current time method to utils
* update calculation
* add tests
* return 0 if stages not running
* address review comments
---
ballista/core/src/utils.rs | 8 ++
ballista/scheduler/src/api/handlers.rs | 162 +++++++++++++++++++++++++++++----
2 files changed, 151 insertions(+), 19 deletions(-)
diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs
index acfc78e46..68e98ea3c 100644
--- a/ballista/core/src/utils.rs
+++ b/ballista/core/src/utils.rs
@@ -349,6 +349,14 @@ pub fn get_time_before(interval_seconds: u64) -> u64 {
.as_secs()
}
+/// current time since UNIX EPOCH
+pub fn get_current_time() -> u128 {
+ SystemTime::now()
+ .duration_since(UNIX_EPOCH)
+ .unwrap()
+ .as_millis()
+}
+
#[cfg(test)]
mod tests {
use super::*;
diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs
index 596f4d09b..87eadf43d 100644
--- a/ballista/scheduler/src/api/handlers.rs
+++ b/ballista/scheduler/src/api/handlers.rs
@@ -13,6 +13,7 @@
use crate::scheduler_server::event::QueryStageSchedulerEvent;
use crate::state::execution_graph::ExecutionStage;
use crate::state::execution_graph_dot::ExecutionGraphDot;
+use crate::state::execution_stage::TaskInfo;
use crate::{api::SchedulerErrorResponse, scheduler_server::SchedulerServer};
use axum::extract::Query;
use axum::{
@@ -28,10 +29,11 @@ use ballista_core::serde::protobuf::{
use ballista_core::serde::scheduler::{
ExecutorOperatingSystemSpecification, ExecutorSpecification,
};
+use ballista_core::utils::get_current_time;
use datafusion::DATAFUSION_VERSION;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::physical_plan::displayable;
-use datafusion::physical_plan::metrics::{MetricValue, MetricsSet, Time};
+use datafusion::physical_plan::metrics::{MetricsSet, Time};
use datafusion_proto::logical_plan::AsLogicalPlan;
use datafusion_proto::physical_plan::AsExecutionPlan;
#[cfg(feature = "graphviz-support")]
@@ -539,11 +541,9 @@ pub async fn get_query_stages<
.as_ref()
.map(|m| get_combined_count(m.as_slice(), "output_rows"))
.unwrap_or(0);
- summary.elapsed_compute = running_stage
- .stage_metrics
- .as_ref()
- .map(|m| get_elapsed_compute_nanos(m.as_slice()))
- .unwrap_or_default();
+ summary.elapsed_compute = get_running_stage_time(&running_stage
+ .task_infos, get_current_time())
+ ;
summary.tasks = running_stage
.task_infos
.iter()
@@ -595,7 +595,7 @@ pub async fn get_query_stages<
"output_rows",
);
summary.elapsed_compute =
- get_elapsed_compute_nanos(&completed_stage.stage_metrics);
+ get_finished_stage_time(&completed_stage.task_infos);
summary.tasks = completed_stage
.task_infos
@@ -722,19 +722,44 @@ fn format_job_status(status: &Option<Status>, elapsed_ms: u64) -> (String, Strin
}
}
-fn get_elapsed_compute_nanos(metrics: &[MetricsSet]) -> String {
- let nanos: usize = metrics
+fn get_running_stage_time(task_infos: &[Option<TaskInfo>], current_time: u128) -> String {
+ let min_start = task_infos
.iter()
- .flat_map(|vec| {
- vec.iter().map(|metric| match metric.as_ref().value() {
- MetricValue::ElapsedCompute(time) => time.value(),
- _ => 0,
- })
- })
- .sum();
- let t = Time::new();
- t.add_duration(Duration::from_nanos(nanos as u64));
- t.to_string()
+ .flat_map(|t| t.as_ref().map(|t| t.start_exec_time))
+ .filter(|t| *t > 0)
+ .min();
+
+ match (min_start, current_time) {
+ (Some(start), end) if end > start => {
+ let t = Time::new();
+ t.add_duration(Duration::from_millis((end - start) as u64));
+ t.to_string()
+ }
+ _ => "0".to_string(),
+ }
+}
+
+fn get_finished_stage_time(task_infos: &[TaskInfo]) -> String {
+ let min_start = task_infos
+ .iter()
+ .map(|t| t.start_exec_time)
+ .filter(|t| *t > 0)
+ .min();
+
+ let max_end = task_infos
+ .iter()
+ .map(|t| t.end_exec_time)
+ .filter(|t| *t > 0)
+ .max();
+
+ match (min_start, max_end) {
+ (Some(start), Some(end)) if end > start => {
+ let t = Time::new();
+ t.add_duration(Duration::from_millis((end - start) as u64));
+ t.to_string()
+ }
+ _ => "0".to_string(),
+ }
}
fn get_partition_counts(metrics: &[MetricsSet], partition_id: usize) -> (usize, usize) {
@@ -878,3 +903,102 @@ pub async fn get_scheduler_metrics<
.unwrap(),
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::state::execution_stage::TaskInfo;
+ use ballista_core::serde::protobuf::task_status;
+
+ fn make_task_info(start: u128, end: u128) -> TaskInfo {
+ TaskInfo {
+ task_id: 0,
+ scheduled_time: 0,
+ launch_time: 0,
+ start_exec_time: start,
+ end_exec_time: end,
+ finish_time: 0,
+ task_status: task_status::Status::Running(Default::default()),
+ }
+ }
+
+ // --- get_finished_stage_time ---
+
+ #[test]
+ fn test_finished_empty_slice_returns_zero() {
+ assert_eq!(get_finished_stage_time(&[]), "0");
+ }
+
+ #[test]
+ fn test_finished_all_zero_timestamps_returns_zero() {
+ let tasks = vec![make_task_info(0, 0), make_task_info(0, 0)];
+ assert_eq!(get_finished_stage_time(&tasks), "0");
+ }
+
+ #[test]
+ fn test_finished_single_task_elapsed() {
+ // 600 - 100 = 500 ms → "500.00ms"
+ let tasks = vec![make_task_info(100, 600)];
+ assert_eq!(get_finished_stage_time(&tasks), "500.00ms");
+ }
+
+ #[test]
+ fn test_finished_picks_earliest_start_and_latest_end() {
+ // min start = 100, max end = 900 → 800 ms
+ let tasks = vec![
+ make_task_info(100, 500),
+ make_task_info(200, 900),
+ make_task_info(300, 700),
+ ];
+ assert_eq!(get_finished_stage_time(&tasks), "800.00ms");
+ }
+
+ #[test]
+ fn test_finished_end_before_start_returns_zero() {
+ let tasks = vec![make_task_info(900, 100)];
+ assert_eq!(get_finished_stage_time(&tasks), "0");
+ }
+
+ // --- get_running_stage_time ---
+
+ #[test]
+ fn test_running_empty_slice_returns_zero() {
+ assert_eq!(get_running_stage_time(&[], 1000), "0");
+ }
+
+ #[test]
+ fn test_running_all_none_returns_zero() {
+ let tasks: Vec<Option<TaskInfo>> = vec![None, None];
+ assert_eq!(get_running_stage_time(&tasks, 1000), "0");
+ }
+
+ #[test]
+ fn test_running_future_start_returns_zero() {
+ // start_exec_time beyond current time → elapsed clamped to 0
+ let tasks = vec![Some(make_task_info(u128::MAX, 0))];
+ assert_eq!(get_running_stage_time(&tasks, 1000), "0");
+ }
+
+ #[test]
+ fn test_running_past_start_returns_nonzero() {
+ let now = 4_000;
+ let start = 1_000;
+ let tasks = vec![Some(make_task_info(start, 0))];
+ assert_eq!(get_running_stage_time(&tasks, now), "3.00s");
+ }
+
+ #[test]
+ fn test_running_mixed_some_none_uses_earliest_some() {
+ let now = 3_000;
+ let earlier = 1_000;
+ let later = 2_000;
+ let tasks = vec![
+ None,
+ Some(make_task_info(later, 0)),
+ Some(make_task_info(earlier, 0)),
+ None,
+ ];
+ let result = get_running_stage_time(&tasks, now);
+ assert_eq!(result, "2.00s");
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org