This is an automated email from the ASF dual-hosted git repository.
andygrove 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 91705602c test: centralize physical plan snapshots (#2297)
91705602c is described below
commit 91705602c9e7aa1b58587905cb73d49371046b97
Author: Floze <[email protected]>
AuthorDate: Sat Aug 15 19:16:30 2026 +0400
test: centralize physical plan snapshots (#2297)
---
ballista/core/src/lib.rs | 16 ++++++++++++++++
ballista/scheduler/src/planner.rs | 15 ++++++---------
ballista/scheduler/src/state/aqe/mod.rs | 12 ------------
.../scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs | 2 +-
.../src/state/aqe/optimizer_rule/distributed_exchange.rs | 2 +-
.../src/state/aqe/optimizer_rule/join_selection.rs | 2 +-
ballista/scheduler/src/state/aqe/test/alter_stages.rs | 2 +-
ballista/scheduler/src/state/aqe/test/coalesce_rule.rs | 2 +-
ballista/scheduler/src/state/aqe/test/join_selection.rs | 6 +++---
ballista/scheduler/src/state/aqe/test/plan_to_stages.rs | 2 +-
10 files changed, 31 insertions(+), 30 deletions(-)
diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs
index feb3f0586..244361eee 100644
--- a/ballista/core/src/lib.rs
+++ b/ballista/core/src/lib.rs
@@ -46,6 +46,22 @@ pub fn print_version() {
println!("Ballista version: {BALLISTA_VERSION}")
}
+/// Asserts an indented physical plan against an inline snapshot.
+///
+/// Keeping plan assertions behind one macro makes plan changes easier to
+/// review and update consistently across Ballista crates.
+#[macro_export]
+macro_rules! assert_plan {
+ ($plan:expr, @ $expected_lines:literal $(,)?) => {{
+ let plan = datafusion::physical_plan::displayable($plan)
+ .indent(true)
+ .to_string();
+ let actual_lines = plan.trim();
+
+ insta::assert_snapshot!(actual_lines, @ $expected_lines);
+ }};
+}
+
/// Client utilities for connecting to Ballista schedulers.
pub mod client;
/// Connection pool for reusing `BallistaClient` instances across requests.
diff --git a/ballista/scheduler/src/planner.rs
b/ballista/scheduler/src/planner.rs
index 14453294e..bb3c80c60 100644
--- a/ballista/scheduler/src/planner.rs
+++ b/ballista/scheduler/src/planner.rs
@@ -928,9 +928,9 @@ pub(crate) fn create_shuffle_writer_with_config(
#[cfg(test)]
mod test {
use super::{can_stay_inline, holds_a_broadcast};
- use crate::assert_plan;
use crate::planner::{DefaultDistributedPlanner, DistributedPlanner};
use crate::test_utils::{datafusion_test_context, scan_with_file_groups};
+ use ballista_core::assert_plan;
use ballista_core::error::BallistaError;
use ballista_core::execution_plans::{SortShuffleWriterExec,
UnresolvedShuffleExec};
use ballista_core::serde::BallistaCodec;
@@ -968,10 +968,10 @@ mod test {
let rewritten = make_empty_exec_serde_safe(plan).unwrap();
- assert_eq!(
- displayable(rewritten.as_ref()).indent(false).to_string(),
- "RepartitionExec: partitioning=RoundRobinBatch(4),
input_partitions=1\n EmptyExec\n"
- );
+ assert_plan!(rewritten.as_ref(), @r"
+ RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+ EmptyExec
+ ");
}
#[test]
@@ -985,10 +985,7 @@ mod test {
let rewritten = make_empty_exec_serde_safe(plan).unwrap();
- assert_eq!(
- displayable(rewritten.as_ref()).indent(false).to_string(),
- "EmptyExec\n"
- );
+ assert_plan!(rewritten.as_ref(), @"EmptyExec");
}
/// A plain scan branch restricts away cleanly, so it can stay inline in
diff --git a/ballista/scheduler/src/state/aqe/mod.rs
b/ballista/scheduler/src/state/aqe/mod.rs
index d06b19d3e..5a1764552 100644
--- a/ballista/scheduler/src/state/aqe/mod.rs
+++ b/ballista/scheduler/src/state/aqe/mod.rs
@@ -1524,15 +1524,3 @@ fn scalar_to_f64(sv: &ScalarValue) ->
datafusion::common::Result<f64> {
),
}
}
-
-/// Checks is the plan same as expected string representation
-#[cfg(test)]
-#[macro_export]
-macro_rules! assert_plan {
- ($PLAN: expr, @ $EXPECTED_LINES: literal $(,)?) => {
- let plan =
datafusion::physical_plan::displayable($PLAN).indent(true).to_string();
- let actual_lines = plan.trim();
-
- insta::assert_snapshot!(actual_lines, @ $EXPECTED_LINES);
- };
-}
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
index bef85f3fc..f7ca5d8f4 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
@@ -109,7 +109,7 @@ impl PhysicalOptimizerRule for ChaosCreatingRule {
#[cfg(test)]
mod tests {
use super::*;
- use crate::assert_plan;
+ use ballista_core::assert_plan;
use ballista_core::config::BallistaConfig;
use ballista_core::execution_plans::ChaosExec;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
diff --git
a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
index 33084d228..822d82f28 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs
@@ -319,8 +319,8 @@ fn nearest_exchange_status(plan: &Arc<dyn ExecutionPlan>)
-> ExchangeStatus {
#[cfg(test)]
mod tests {
use super::*;
- use crate::assert_plan;
use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec,
ExchangeExec};
+ use ballista_core::assert_plan;
use ballista_core::execution_plans::{
RuntimeStatsExec, UnorderedRangeRepartitionExec,
};
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
index 78b8fa1c1..ce4f79146 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs
@@ -365,7 +365,7 @@ mod tests {
use super::*;
use std::sync::Arc;
- use crate::assert_plan;
+ use ballista_core::assert_plan;
use ballista_core::config::BallistaConfig;
use datafusion::{
arrow::{
diff --git a/ballista/scheduler/src/state/aqe/test/alter_stages.rs
b/ballista/scheduler/src/state/aqe/test/alter_stages.rs
index bb31ab7bc..387f1c58b 100644
--- a/ballista/scheduler/src/state/aqe/test/alter_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/alter_stages.rs
@@ -15,11 +15,11 @@
// specific language governing permissions and limitations
// under the License.
-use crate::assert_plan;
use crate::state::aqe::planner::AdaptivePlanner;
use crate::state::aqe::test::{
mock_batch, mock_context, mock_partitions_with_statistics_no_data,
};
+use ballista_core::assert_plan;
use ballista_core::serde::scheduler::{
ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
PartitionId, PartitionLocation, PartitionStats,
diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
index c51cff0a1..a3da8bb8b 100644
--- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
+++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs
@@ -24,9 +24,9 @@
//! `coalesce_target_partition_bytes` so the bin-pack outcome is hand-traceable
//! against `split_size_list_by_target_size`.
-use crate::assert_plan;
use crate::state::aqe::planner::AdaptivePlanner;
use crate::state::aqe::test::{mock_batch, mock_schema};
+use ballista_core::assert_plan;
use ballista_core::extension::SessionConfigExt;
use ballista_core::serde::scheduler::{
ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
diff --git a/ballista/scheduler/src/state/aqe/test/join_selection.rs
b/ballista/scheduler/src/state/aqe/test/join_selection.rs
index 9bf234778..6243d7010 100644
--- a/ballista/scheduler/src/state/aqe/test/join_selection.rs
+++ b/ballista/scheduler/src/state/aqe/test/join_selection.rs
@@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.
-use crate::{
- assert_plan,
- state::aqe::{planner::AdaptivePlanner,
test::mock_partitions_with_statistics},
+use crate::state::aqe::{
+ planner::AdaptivePlanner, test::mock_partitions_with_statistics,
};
+use ballista_core::assert_plan;
use ballista_core::extension::SessionConfigExt;
use datafusion::{
arrow::{
diff --git a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
index 5290ff26d..3a18933e7 100644
--- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
+++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs
@@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
-use crate::assert_plan;
use crate::state::aqe::AdaptiveExecutionGraph;
use crate::state::aqe::execution_plan::ExchangeExec;
use crate::state::aqe::planner::AdaptivePlanner;
@@ -24,6 +23,7 @@ use crate::state::aqe::test::{
};
use crate::state::execution_graph::ExecutionGraph;
use ballista_core::JobId;
+use ballista_core::assert_plan;
use ballista_core::execution_plans::SortShuffleWriterExec;
use ballista_core::serde::protobuf::job_status::Status;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
