From 46338d140868929bb3af531ab5219459083e9ab1 Mon Sep 17 00:00:00 2001 From: Hoang Pham Date: Sun, 28 Dec 2025 23:14:12 +0700 Subject: [PATCH 1/2] add more metrics Signed-off-by: Hoang Pham --- ballista/client/tests/context_checks.rs | 39 ++++++++++++++ .../src/execution_plans/distributed_query.rs | 52 +++++++++++++++++-- 2 files changed, 87 insertions(+), 4 deletions(-) diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index deba3bc33e..7bc1d00bd2 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -123,6 +123,31 @@ mod supported { > 0 ); + // Verify timing metrics + let job_execution_time = metrics.sum_by_name("job_execution_time_ms").unwrap(); + assert!( + job_execution_time.as_usize() > 0, + "job_execution_time_ms should be greater than 0" + ); + + let scheduling_time = metrics.sum_by_name("job_scheduling_in_ms").unwrap(); + assert!( + scheduling_time.as_usize() > 0, + "job_scheduling_in_ms should be non-negative" + ); + + let total_time = metrics.sum_by_name("total_query_time_ms").unwrap(); + assert!( + total_time.as_usize() > 0, + "total_query_time_ms should be greater than 0" + ); + + // Total time should be at least as long as execution time + assert!( + total_time.as_usize() >= job_execution_time.as_usize(), + "total_query_time_ms should be >= job_execution_time_ms" + ); + Ok(()) } @@ -182,6 +207,20 @@ mod supported { assert_batches_eq!(expected, &result); + // Verify timing metrics + let metrics = plan.metrics().unwrap(); + let job_execution_time = metrics.sum_by_name("job_execution_time_ms").unwrap(); + assert!( + job_execution_time.as_usize() > 0, + "job_execution_time_ms should be greater than 0" + ); + + let total_time = metrics.sum_by_name("total_query_time_ms").unwrap(); + assert!( + total_time.as_usize() >= job_execution_time.as_usize(), + "total_query_time_ms should be >= job_execution_time_ms" + ); + Ok(()) } diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 4336b1fc3b..354e05df9a 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -71,8 +71,12 @@ pub struct DistributedQueryExec { /// Plan properties properties: PlanProperties, /// Execution metrics, currently exposes: - /// - row count - /// - transferred_bytes + /// - output_rows: Total number of rows returned + /// - transferred_bytes: Total bytes transferred from executors + /// - job_execution_time_ms: Time spent executing on the cluster (server-side) + /// - job_scheduling_in_ms: Time from query submission to job start (includes queue time) + /// - total_query_time_ms: Total end-to-end query latency from client perspective + /// - data_transfer_time_ms: Time spent fetching results (not yet implemented) metrics: ExecutionPlanMetricsSet, } @@ -234,6 +238,7 @@ impl ExecutionPlan for DistributedQueryExec { let metric_row_count = MetricBuilder::new(&self.metrics).output_rows(partition); let metric_total_bytes = MetricBuilder::new(&self.metrics).counter("transferred_bytes", partition); + let stream = futures::stream::once( execute_query( self.scheduler_url.clone(), @@ -241,6 +246,8 @@ impl ExecutionPlan for DistributedQueryExec { query, self.config.default_grpc_client_max_message_size(), GrpcClientConfig::from(&self.config), + Arc::new(self.metrics.clone()), + partition, ) .map_err(|e| ArrowError::ExternalError(Box::new(e))), ) @@ -278,7 +285,12 @@ async fn execute_query( query: ExecuteQueryParams, max_message_size: usize, grpc_config: GrpcClientConfig, + metrics: Arc, + partition: usize, ) -> Result> + Send> { + // Capture query submission time for total_query_time_ms + let query_start_time = std::time::Instant::now(); + info!("Connecting to Ballista scheduler at {scheduler_url}"); // TODO reuse the scheduler to avoid connecting to the Ballista scheduler again and again let connection = create_grpc_client_connection(scheduler_url, &grpc_config) @@ -289,6 +301,8 @@ async fn execute_query( .max_encoding_message_size(max_message_size) .max_decoding_message_size(max_message_size); + let query_submit_time = std::time::Instant::now(); + let query_result = scheduler .execute_query(query) .await @@ -356,10 +370,40 @@ async fn execute_query( partition_location, .. })) => { - let duration = ended_at.saturating_sub(started_at); - let duration = Duration::from_millis(duration); + // Calculate job execution time (server-side execution) + let job_execution_ms = ended_at.saturating_sub(started_at); + let duration = Duration::from_millis(job_execution_ms); info!("Job {job_id} finished executing in {duration:?} "); + + // Calculate scheduling time (client-perceived time from query submission to job start) + // This includes network latency and actual queue time + let scheduling_elapsed = query_submit_time.elapsed(); + let scheduling_ms = scheduling_elapsed + .as_millis() + .saturating_sub(job_execution_ms as u128); + + // Calculate total query time (end-to-end from client perspective) + let total_elapsed = query_start_time.elapsed(); + let total_ms = total_elapsed.as_millis(); + + // Set timing metrics + let metric_job_execution = MetricBuilder::new(&metrics) + .gauge("job_execution_time_ms", partition); + metric_job_execution.set(job_execution_ms as usize); + + let metric_scheduling = + MetricBuilder::new(&metrics).gauge("job_scheduling_in_ms", partition); + metric_scheduling.set(scheduling_ms as usize); + + let metric_total_time = + MetricBuilder::new(&metrics).gauge("total_query_time_ms", partition); + metric_total_time.set(total_ms as usize); + + // Note: data_transfer_time_ms is not set here because partition fetching + // happens lazily when the stream is consumed, not during execute_query. + // This could be added in a future enhancement by wrapping the stream. + let streams = partition_location.into_iter().map(move |partition| { let f = fetch_partition(partition, max_message_size, true) .map_err(|e| ArrowError::ExternalError(Box::new(e))); From 366673a1cbd0332bcc3d336eb59531a86da48725 Mon Sep 17 00:00:00 2001 From: Hoang Pham Date: Tue, 30 Dec 2025 13:06:02 +0700 Subject: [PATCH 2/2] use queue_at Signed-off-by: Hoang Pham --- .../core/src/execution_plans/distributed_query.rs | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 354e05df9a..de35212658 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -75,8 +75,8 @@ pub struct DistributedQueryExec { /// - transferred_bytes: Total bytes transferred from executors /// - job_execution_time_ms: Time spent executing on the cluster (server-side) /// - job_scheduling_in_ms: Time from query submission to job start (includes queue time) - /// - total_query_time_ms: Total end-to-end query latency from client perspective - /// - data_transfer_time_ms: Time spent fetching results (not yet implemented) + /// - job_execution_time_ms: Time spent executing on the cluster (ended_at - started_at) + /// - job_scheduling_in_ms: Time job waited in scheduler queue (started_at - queued_at) metrics: ExecutionPlanMetricsSet, } @@ -301,8 +301,6 @@ async fn execute_query( .max_encoding_message_size(max_message_size) .max_decoding_message_size(max_message_size); - let query_submit_time = std::time::Instant::now(); - let query_result = scheduler .execute_query(query) .await @@ -365,6 +363,7 @@ async fn execute_query( break Err(DataFusionError::Execution(msg)); } Some(job_status::Status::Successful(SuccessfulJob { + queued_at, started_at, ended_at, partition_location, @@ -376,12 +375,9 @@ async fn execute_query( info!("Job {job_id} finished executing in {duration:?} "); - // Calculate scheduling time (client-perceived time from query submission to job start) + // Calculate scheduling time (server-side queue time) // This includes network latency and actual queue time - let scheduling_elapsed = query_submit_time.elapsed(); - let scheduling_ms = scheduling_elapsed - .as_millis() - .saturating_sub(job_execution_ms as u128); + let scheduling_ms = started_at.saturating_sub(queued_at); // Calculate total query time (end-to-end from client perspective) let total_elapsed = query_start_time.elapsed();