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 cd53c729 ballista client will collect few metrics (#1251)
cd53c729 is described below
commit cd53c7291a151d739ff8f07d3c53365eed8398a4
Author: Marko Milenković <[email protected]>
AuthorDate: Tue Apr 29 09:52:45 2025 +0100
ballista client will collect few metrics (#1251)
---
ballista/client/tests/context_checks.rs | 117 +++++++++++++++++++++
.../core/src/execution_plans/distributed_query.rs | 62 ++++++++---
ballista/scheduler/src/state/execution_graph.rs | 9 +-
3 files changed, 170 insertions(+), 18 deletions(-)
diff --git a/ballista/client/tests/context_checks.rs
b/ballista/client/tests/context_checks.rs
index d27350f7..f60559ac 100644
--- a/ballista/client/tests/context_checks.rs
+++ b/ballista/client/tests/context_checks.rs
@@ -21,6 +21,7 @@ mod supported {
use crate::common::{remote_context, standalone_context};
use ballista_core::config::BallistaConfig;
+ use datafusion::physical_plan::collect;
use datafusion::prelude::*;
use datafusion::{assert_batches_eq, prelude::SessionContext};
use rstest::*;
@@ -66,6 +67,122 @@ mod supported {
Ok(())
}
+
+ // tests if client will collect statistics for
+ // collect/show operation
+ #[rstest]
+ #[case::standalone(standalone_context())]
+ #[case::remote(remote_context())]
+ #[tokio::test]
+ async fn should_collect_client_statistics_for_show(
+ #[future(awt)]
+ #[case]
+ ctx: SessionContext,
+ test_data: String,
+ ) -> datafusion::error::Result<()> {
+ ctx.register_parquet(
+ "test",
+ &format!("{test_data}/alltypes_plain.parquet"),
+ Default::default(),
+ )
+ .await?;
+
+ let plan = ctx
+ .sql("select string_col, timestamp_col from test where id > 4")
+ .await?
+ .create_physical_plan()
+ .await?;
+
+ let result = collect(plan.clone(), ctx.task_ctx()).await?;
+
+ let expected = [
+ "+------------+---------------------+",
+ "| string_col | timestamp_col |",
+ "+------------+---------------------+",
+ "| 31 | 2009-03-01T00:01:00 |",
+ "| 30 | 2009-04-01T00:00:00 |",
+ "| 31 | 2009-04-01T00:01:00 |",
+ "+------------+---------------------+",
+ ];
+
+ assert_batches_eq!(expected, &result);
+
+ let metrics = plan.metrics().unwrap();
+ let rows = metrics.output_rows().unwrap();
+
+ assert_eq!(3, rows);
+ assert!(
+ plan.metrics()
+ .unwrap()
+ .sum_by_name("transferred_bytes")
+ .unwrap()
+ .as_usize()
+ > 0
+ );
+
+ Ok(())
+ }
+
+ // tests if client will collect statistics for
+ // insert operation
+ #[rstest]
+ #[case::standalone(standalone_context())]
+ #[case::remote(remote_context())]
+ #[tokio::test]
+ #[cfg(not(windows))] // test is failing at windows, can't debug it
+ async fn should_collect_client_statistics_for_insert(
+ #[future(awt)]
+ #[case]
+ ctx: SessionContext,
+ test_data: String,
+ ) -> datafusion::error::Result<()> {
+ ctx.register_parquet(
+ "test",
+ &format!("{test_data}/alltypes_plain.parquet"),
+ Default::default(),
+ )
+ .await
+ .unwrap();
+ let write_dir = tempfile::tempdir().expect("temporary directory to be
created");
+ let write_dir_path = write_dir
+ .path()
+ .to_str()
+ .expect("path to be converted to str");
+
+ ctx.sql(&format!("CREATE EXTERNAL TABLE written_table (id INTEGER,
string_col STRING, timestamp_col BIGINT) STORED AS PARQUET LOCATION '{}'" ,
write_dir_path)).await.unwrap().show().await.unwrap();
+
+ let plan = ctx
+ .sql("INSERT INTO written_table select id, string_col,
timestamp_col from test")
+ .await?
+ .create_physical_plan()
+ .await?;
+
+ let result = collect(plan.clone(), ctx.task_ctx()).await.unwrap();
+
+ // INSERT operation should return only single row
+ assert_eq!(1, plan.metrics().unwrap().output_rows().unwrap());
+ assert!(
+ plan.metrics()
+ .unwrap()
+ .sum_by_name("transferred_bytes")
+ .unwrap()
+ .as_usize()
+ > 0
+ );
+
+ let expected = [
+ "+-------+",
+ "| count |",
+ "+-------+",
+ "| 8 |",
+ "+-------+",
+ ];
+
+ assert_batches_eq!(expected, &result);
+
+ Ok(())
+ }
+
#[rstest]
#[case::standalone(standalone_context())]
#[case::remote(remote_context())]
diff --git a/ballista/core/src/execution_plans/distributed_query.rs
b/ballista/core/src/execution_plans/distributed_query.rs
index 59503814..7ccb2727 100644
--- a/ballista/core/src/execution_plans/distributed_query.rs
+++ b/ballista/core/src/execution_plans/distributed_query.rs
@@ -17,6 +17,7 @@
use crate::client::BallistaClient;
use crate::config::BallistaConfig;
+use crate::serde::protobuf::SuccessfulJob;
use crate::serde::protobuf::{
execute_query_params::Query, execute_query_result, job_status,
scheduler_grpc_client::SchedulerGrpcClient, ExecuteQueryParams,
GetJobStatusParams,
@@ -30,6 +31,9 @@ use datafusion::error::{DataFusionError, Result};
use datafusion::execution::context::TaskContext;
use datafusion::logical_expr::LogicalPlan;
use datafusion::physical_expr::EquivalenceProperties;
+use datafusion::physical_plan::metrics::{
+ ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
+};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
@@ -64,7 +68,12 @@ pub struct DistributedQueryExec<T: 'static + AsLogicalPlan> {
plan_repr: PhantomData<T>,
/// Session id
session_id: String,
+ /// Plan properties
properties: PlanProperties,
+ /// Execution metrics, currently exposes:
+ /// - row count
+ /// - transferred_bytes
+ metrics: ExecutionPlanMetricsSet,
}
impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
@@ -83,6 +92,7 @@ impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
plan_repr: PhantomData,
session_id,
properties,
+ metrics: ExecutionPlanMetricsSet::new(),
}
}
@@ -102,6 +112,7 @@ impl<T: 'static + AsLogicalPlan> DistributedQueryExec<T> {
plan_repr: PhantomData,
session_id,
properties,
+ metrics: ExecutionPlanMetricsSet::new(),
}
}
@@ -168,6 +179,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for
DistributedQueryExec<T> {
properties: Self::compute_properties(
self.plan.schema().as_ref().clone().into(),
),
+ metrics: ExecutionPlanMetricsSet::new(),
}))
}
@@ -209,6 +221,9 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for
DistributedQueryExec<T> {
session_id: Some(self.session_id.clone()),
};
+ 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(),
@@ -218,7 +233,17 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for
DistributedQueryExec<T> {
)
.map_err(|e| ArrowError::ExternalError(Box::new(e))),
)
- .try_flatten();
+ .try_flatten()
+ .inspect(move |batch| {
+ metric_total_bytes.add(
+ batch
+ .as_ref()
+ .map(|b| b.get_array_memory_size())
+ .unwrap_or(0),
+ );
+
+ metric_row_count.add(batch.as_ref().map(|b|
b.num_rows()).unwrap_or(0));
+ });
let schema = self.schema();
Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
@@ -230,6 +255,10 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for
DistributedQueryExec<T> {
// This implies that we cannot infer the statistics at this stage.
Ok(Statistics::new_unknown(&self.schema()))
}
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
}
async fn execute_query(
@@ -285,14 +314,14 @@ async fn execute_query(
match status {
None => {
if has_status_change {
- info!("Job {} still in initialization ...", job_id);
+ info!("Job {} is in initialization ...", job_id);
}
wait_future.await;
prev_status = status;
}
Some(job_status::Status::Queued(_)) => {
if has_status_change {
- info!("Job {} still queued...", job_id);
+ info!("Job {} is queued...", job_id);
}
wait_future.await;
prev_status = status;
@@ -309,17 +338,22 @@ async fn execute_query(
error!("{}", msg);
break Err(DataFusionError::Execution(msg));
}
- Some(job_status::Status::Successful(successful)) => {
- let streams =
- successful
- .partition_location
- .into_iter()
- .map(move |partition| {
- let f = fetch_partition(partition,
max_message_size)
- .map_err(|e|
ArrowError::ExternalError(Box::new(e)));
-
- futures::stream::once(f).try_flatten()
- });
+ Some(job_status::Status::Successful(SuccessfulJob {
+ started_at,
+ ended_at,
+ partition_location,
+ ..
+ })) => {
+ let duration = ended_at.saturating_sub(started_at);
+ let duration = Duration::from_millis(duration);
+
+ info!("Job {} finished executing in {:?} ", job_id, duration);
+ let streams = partition_location.into_iter().map(move
|partition| {
+ let f = fetch_partition(partition, max_message_size)
+ .map_err(|e| ArrowError::ExternalError(Box::new(e)));
+
+ futures::stream::once(f).try_flatten()
+ });
break Ok(futures::stream::iter(streams).flatten());
}
diff --git a/ballista/scheduler/src/state/execution_graph.rs
b/ballista/scheduler/src/state/execution_graph.rs
index 9023a7ec..7114ba6e 100644
--- a/ballista/scheduler/src/state/execution_graph.rs
+++ b/ballista/scheduler/src/state/execution_graph.rs
@@ -1293,6 +1293,11 @@ impl ExecutionGraph {
.map(|l| l.try_into())
.collect::<Result<Vec<_>>>()?;
+ self.end_time = SystemTime::now()
+ .duration_since(UNIX_EPOCH)
+ .unwrap()
+ .as_millis() as u64;
+
self.status = JobStatus {
job_id: self.job_id.clone(),
job_name: self.job_name.clone(),
@@ -1304,10 +1309,6 @@ impl ExecutionGraph {
ended_at: self.end_time,
})),
};
- self.end_time = SystemTime::now()
- .duration_since(UNIX_EPOCH)
- .unwrap()
- .as_millis() as u64;
Ok(())
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
