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/06/02 17:22:53 UTC
(datafusion-ballista) branch main updated: feat(aqe): chaos monkey testing for ballista execution (#1802)
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 73450f555 feat(aqe): chaos monkey testing for ballista execution (#1802)
73450f555 is described below
commit 73450f555df823c389c182fbc0083a7123f3ad08
Author: Marko Milenković <mi...@users.noreply.github.com>
AuthorDate: Tue Jun 2 18:22:48 2026 +0100
feat(aqe): chaos monkey testing for ballista execution (#1802)
* Add chaos execution plan and physical optimizer rule
* Update ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* Update ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* Update ballista/core/src/execution_plans/chaos_exec.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* Update ballista/core/src/config.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* Update ballista/core/src/config.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* Update ballista/core/src/serde/mod.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* Update ballista/core/src/execution_plans/chaos_exec.rs
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
* seed is predetermined at the creation time
that way its visible in the plan and can be used
to regenerate same plans if needed
---------
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
---
ballista/core/proto/ballista.proto | 7 +
ballista/core/src/config.rs | 67 +++++
ballista/core/src/execution_plans/chaos_exec.rs | 314 +++++++++++++++++++++
ballista/core/src/execution_plans/mod.rs | 2 +
ballista/core/src/serde/generated/ballista.rs | 13 +-
ballista/core/src/serde/mod.rs | 35 ++-
.../src/state/aqe/optimizer_rule/chaos_exec.rs | 204 +++++++++++++
.../scheduler/src/state/aqe/optimizer_rule/mod.rs | 1 +
ballista/scheduler/src/state/aqe/planner.rs | 6 +-
9 files changed, 646 insertions(+), 3 deletions(-)
diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto
index cb209b55e..2154b1a17 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -50,9 +50,16 @@ message BallistaPhysicalPlanNode {
ShuffleReaderExecNode shuffle_reader = 2;
UnresolvedShuffleExecNode unresolved_shuffle = 3;
SortShuffleWriterExecNode sort_shuffle_writer = 4;
+ ChaosExecNode chaos_exec = 5;
}
}
+message ChaosExecNode {
+ double failure_probability = 1;
+ string fault_type = 2;
+ uint64 seed = 3;
+}
+
message ShuffleWriterExecNode {
//TODO it seems redundant to provide job and stage id here since we also have them
// in the TaskDefinition that wraps this plan
diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs
index abb8befc0..8f87d5384 100644
--- a/ballista/core/src/config.rs
+++ b/ballista/core/src/config.rs
@@ -104,6 +104,19 @@ pub const BALLISTA_COALESCE_MERGED_PARTITION_FACTOR: &str =
/// When disabled, `SelectJoinRule` is a no-op and `DynamicJoinSelectionExec`
/// nodes are left in the plan (which will fail at execution — disable only for debugging).
pub const BALLISTA_ADAPTIVE_JOIN_ENABLED: &str = "ballista.planner.adaptive_join.enabled";
+/// Configuration key to enable chaos-monkey execution injection for robustness testing.
+pub const BALLISTA_CHAOS_EXECUTION_ENABLED: &str =
+ "ballista.testing.chaos_execution.enabled";
+/// Configuration key for the per-node failure probability used by chaos-monkey execution.
+pub const BALLISTA_CHAOS_EXECUTION_PROBABILITY: &str =
+ "ballista.testing.chaos_execution.probability";
+/// Configuration key controlling the fault type injected by chaos-monkey execution.
+/// Valid values: `"transient"` (IoError), `"fatal"` (Execution error), `"panic"`, `"delay"`.
+pub const BALLISTA_CHAOS_EXECUTION_FAULT_TYPE: &str =
+ "ballista.testing.chaos_execution.fault_type";
+/// Configuration key for the optional RNG seed used by chaos-monkey execution.
+/// An empty string (default) means non-deterministic; set to a `u64` value for reproducibility.
+pub const BALLISTA_CHAOS_EXECUTION_SEED: &str = "ballista.testing.chaos_execution.seed";
/// Result type for configuration parsing operations.
pub type ParseResult<T> = result::Result<T, String>;
@@ -233,6 +246,38 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
DataType::Boolean,
Some(true.to_string()),
),
+ ConfigEntry::new(
+ BALLISTA_CHAOS_EXECUTION_ENABLED.to_string(),
+ "Enables chaos-monkey execution injection for robustness testing. \
+ When true, ChaosExec is inserted at a random point in the plan \
+ once per optimize call.".to_string(),
+ DataType::Boolean,
+ Some(false.to_string()),
+ ),
+ ConfigEntry::new(
+ BALLISTA_CHAOS_EXECUTION_PROBABILITY.to_string(),
+ "Failure probability (0.0–1.0) passed to ChaosExec when \
+ chaos execution is enabled.".to_string(),
+ DataType::Float64,
+ Some("0.25".to_string()),
+ ),
+ ConfigEntry::new(
+ BALLISTA_CHAOS_EXECUTION_FAULT_TYPE.to_string(),
+ "Fault type injected by chaos-monkey execution. \
+ \"transient\": IoError (retryable); \"fatal\": Execution error (non-retryable); \
+ \"panic\": panics the task thread; \
+ \"delay\" or \"delay:N\": sleeps N ms per batch with no error (default N=1).".to_string(),
+ DataType::Utf8,
+ Some("transient".to_string()),
+ ),
+ ConfigEntry::new(
+ BALLISTA_CHAOS_EXECUTION_SEED.to_string(),
+ "Optional u64 seed for the chaos RNG. \
+ Empty string (default) means non-deterministic; set to a numeric value \
+ to get reproducible fault injection across runs.".to_string(),
+ DataType::Utf8,
+ Some("".to_string()),
+ ),
];
entries
.into_iter()
@@ -468,6 +513,28 @@ impl BallistaConfig {
self.get_bool_setting(BALLISTA_ADAPTIVE_JOIN_ENABLED)
}
+ /// Returns whether chaos-monkey execution injection is enabled.
+ pub fn chaos_execution_enabled(&self) -> bool {
+ self.get_bool_setting(BALLISTA_CHAOS_EXECUTION_ENABLED)
+ }
+
+ /// Returns the failure probability for chaos-monkey execution (0.0–1.0).
+ pub fn chaos_execution_probability(&self) -> f64 {
+ self.get_float_setting(BALLISTA_CHAOS_EXECUTION_PROBABILITY)
+ }
+
+ /// Returns the fault type injected by chaos-monkey execution (default `"transient"`).
+ pub fn chaos_execution_fault_type(&self) -> String {
+ self.get_string_setting(BALLISTA_CHAOS_EXECUTION_FAULT_TYPE)
+ }
+
+ /// Returns the optional RNG seed for chaos-monkey execution.
+ /// `None` means non-deterministic (thread-local RNG); `Some(n)` gives reproducible runs.
+ pub fn chaos_execution_seed(&self) -> Option<u64> {
+ let s = self.get_string_setting(BALLISTA_CHAOS_EXECUTION_SEED);
+ if s.is_empty() { None } else { s.parse().ok() }
+ }
+
/// Should client employ pull or push job tracking strategy
pub fn client_pull(&self) -> bool {
self.get_bool_setting(BALLISTA_CLIENT_PULL)
diff --git a/ballista/core/src/execution_plans/chaos_exec.rs b/ballista/core/src/execution_plans/chaos_exec.rs
new file mode 100644
index 000000000..fda436ec5
--- /dev/null
+++ b/ballista/core/src/execution_plans/chaos_exec.rs
@@ -0,0 +1,314 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// ChaosExec is a physical execution plan node for robustness/chaos testing.
+// It wraps a single child node, preserving its schema and partitioning, and
+// randomly injects faults according to a configurable failure_probability
+// (in [0.0, 1.0]) and fault_type:
+//
+// "transient" — returns a recoverable IoError on the first batch
+// "fatal" — returns a non-recoverable Execution error on the first batch
+// "panic" — panics on the first batch
+// "delay" — sleeps 1 ms before every batch
+// "delay:N" — sleeps N ms before every batch
+//
+// ChaosExec is inserted into query plans by physical optimizer rule, which
+// probabilistically wraps leaf nodes.
+
+use datafusion::common::{DataFusionError, Result, Statistics, internal_err};
+use datafusion::config::ConfigOptions;
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::execution_plan::CardinalityEffect;
+use datafusion::physical_plan::metrics::MetricsSet;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+ SendableRecordBatchStream,
+};
+use futures::StreamExt;
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+use std::any::Any;
+use std::sync::Arc;
+
+/// Physical execution plan node that randomly injects failures for chaos/robustness testing.
+#[derive(Debug, Clone)]
+pub struct ChaosExec {
+ input: Arc<dyn ExecutionPlan>,
+ // a probability this node will fail when run
+ failure_probability: f64,
+ // controls what kind of fault is injected: "transient", "fatal", "panic", or "delay"
+ fault_type: String,
+ // seed used for execution
+ seed: u64,
+}
+
+impl ChaosExec {
+ /// Creates a new `ChaosExec` wrapping `input`.
+ ///
+ /// `failure_probability` must be in `[0.0, 1.0]`.
+ /// `fault_type` must be one of `"transient"`, `"fatal"`, `"panic"`, `"delay"`, or `"delay:N"`
+ /// where N is a positive integer number of milliseconds.
+ pub fn new(
+ input: Arc<dyn ExecutionPlan>,
+ failure_probability: f64,
+ fault_type: &str,
+ seed: Option<u64>,
+ ) -> Result<Self> {
+ let seed = match seed {
+ Some(seed) => seed,
+ None => rand::random::<u64>(),
+ };
+ if !(0.0..=1.0).contains(&failure_probability) {
+ return internal_err!(
+ "ChaosExec failure_probability must be in [0.0, 1.0], got {failure_probability}"
+ );
+ }
+ match fault_type {
+ "transient" | "fatal" | "panic" | "delay" => {}
+ other if other.starts_with("delay:") => {
+ let ms_str = &other["delay:".len()..];
+ if ms_str.parse::<u64>().is_err() {
+ return internal_err!(
+ "ChaosExec delay suffix must be a positive integer (ms), got \"{ms_str}\""
+ );
+ }
+ }
+ other => {
+ return internal_err!(
+ "ChaosExec fault_type must be one of transient/fatal/panic/delay/delay:N, got {other}"
+ );
+ }
+ }
+
+ Ok(Self {
+ input,
+ failure_probability,
+ fault_type: fault_type.to_string(),
+ seed,
+ })
+ }
+
+ /// Returns the configured RNG seed.
+ pub fn seed(&self) -> u64 {
+ self.seed
+ }
+
+ /// Returns the configured failure probability.
+ pub fn failure_probability(&self) -> f64 {
+ self.failure_probability
+ }
+
+ /// Returns the configured fault type.
+ pub fn fault_type(&self) -> &str {
+ &self.fault_type
+ }
+}
+
+impl ExecutionPlan for ChaosExec {
+ fn name(&self) -> &str {
+ "ChaosExec"
+ }
+
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ self.input.properties()
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![&self.input]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ mut children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if children.len() != 1 {
+ return internal_err!("ChaosExec expected one child, got {}", children.len());
+ }
+ let new_input = children.pop().unwrap();
+ Ok(Arc::new(Self::new(
+ new_input,
+ self.failure_probability,
+ &self.fault_type,
+ Some(self.seed),
+ )?))
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let input_stream = self.input.execute(partition, context)?;
+ let failure_probability = self.failure_probability;
+ let fault_type = self.fault_type.to_lowercase();
+ let schema = self.input.schema();
+
+ // Decide once per partition execution whether to inject a fault on the first batch.
+ // With a seed, mix in the partition index so each partition has an independent but
+ // deterministic sequence; without a seed, fall back to the thread-local RNG.
+ let should_fail = {
+ let mut rng = StdRng::seed_from_u64(self.seed.wrapping_add(partition as u64));
+ rng.random::<f64>() < failure_probability
+ };
+
+ // Wrap the child stream. For error/panic modes, inject on the first batch (idx == 0)
+ // to mirror how real IO failures surface in production. For "delay", sleep every batch.
+ let wrapped = input_stream.enumerate().map(move |(idx, batch_result)| {
+ match fault_type.as_str() {
+ ft if ft.starts_with("delay") => {
+ std::thread::sleep(std::time::Duration::from_millis(parse_delay_ms(ft)));
+ batch_result
+ }
+ "transient" if idx == 0 && should_fail => {
+ let error_msg = format!(
+ "ChaosExec: Injected TRANSIENT FAILURE (recoverable) on partition {partition}"
+ );
+ log::error!("{}",error_msg);
+ Err(DataFusionError::IoError(std::io::Error::other(error_msg)))
+ }
+ "fatal" if idx == 0 && should_fail => {
+ let error_msg = format!(
+ "ChaosExec: Injected FATAL FAILURE on partition {partition} (chaos testing)"
+ );
+ log::error!("{}",error_msg);
+ Err(DataFusionError::Execution(error_msg))
+ }
+ "panic" if idx == 0 && should_fail => {
+ log::error!("ChaosExec: Injected panic on partition {partition}");
+ panic!("ChaosExec: injected PANIC on partition {partition}")
+ }
+ "transient" | "fatal" | "panic" => batch_result,
+ config => {
+ let error_msg = format!(
+ "ChaosExec: wrong config value {config}, will break execution anyway, "
+ );
+ log::error!("{}",error_msg);
+ Err(DataFusionError::Configuration(error_msg))
+ },
+ }
+ });
+
+ Ok(Box::pin(RecordBatchStreamAdapter::new(schema, wrapped)))
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ self.input.metrics()
+ }
+
+ fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
+ self.input.partition_statistics(partition)
+ }
+
+ fn supports_limit_pushdown(&self) -> bool {
+ self.input.supports_limit_pushdown()
+ }
+
+ fn fetch(&self) -> Option<usize> {
+ self.input.fetch()
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ self.input.cardinality_effect()
+ }
+
+ fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
+ let new_input = self.input.with_fetch(limit)?;
+ Some(Arc::new(
+ Self::new(
+ new_input,
+ self.failure_probability,
+ &self.fault_type,
+ Some(self.seed),
+ )
+ .ok()?,
+ ))
+ }
+
+ fn repartitioned(
+ &self,
+ target_partitions: usize,
+ config: &ConfigOptions,
+ ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ if let Some(new_input) = self.input.repartitioned(target_partitions, config)? {
+ Ok(Some(Arc::new(Self::new(
+ new_input,
+ self.failure_probability,
+ &self.fault_type,
+ Some(self.seed),
+ )?)))
+ } else {
+ Ok(None)
+ }
+ }
+
+ fn with_preserve_order(
+ &self,
+ preserve_order: bool,
+ ) -> Option<Arc<dyn ExecutionPlan>> {
+ let new_input = self.input.with_preserve_order(preserve_order)?;
+ Some(Arc::new(
+ Self::new(
+ new_input,
+ self.failure_probability,
+ &self.fault_type,
+ Some(self.seed),
+ )
+ .ok()?,
+ ))
+ }
+}
+
+fn parse_delay_ms(fault_type: &str) -> u64 {
+ fault_type
+ .strip_prefix("delay:")
+ .and_then(|s| s.parse().ok())
+ .unwrap_or(1)
+}
+
+impl DisplayAs for ChaosExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(
+ f,
+ "ChaosExec: failure_probability={}, fault_type={}, seed={}",
+ self.failure_probability, self.fault_type, self.seed
+ )
+ }
+ DisplayFormatType::TreeRender => {
+ writeln!(
+ f,
+ "failure_probability={}, fault_type={}, seed={}",
+ self.failure_probability, self.fault_type, self.seed
+ )
+ }
+ }
+ }
+}
diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs
index ae46fad68..91bfddf45 100644
--- a/ballista/core/src/execution_plans/mod.rs
+++ b/ballista/core/src/execution_plans/mod.rs
@@ -18,6 +18,7 @@
//! This module contains execution plans that are needed to distribute DataFusion's execution plans into
//! several Ballista executors.
+mod chaos_exec;
mod distributed_explain_analyze;
mod distributed_query;
mod shuffle_reader;
@@ -28,6 +29,7 @@ mod unresolved_shuffle;
use std::path::{Path, PathBuf};
+pub use chaos_exec::ChaosExec;
use datafusion::common::exec_err;
pub use distributed_explain_analyze::DistributedExplainAnalyzeExec;
pub use distributed_query::DistributedQueryExec;
diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs
index 774c53ad3..269827145 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode {
pub struct BallistaPhysicalPlanNode {
#[prost(
oneof = "ballista_physical_plan_node::PhysicalPlanType",
- tags = "1, 2, 3, 4"
+ tags = "1, 2, 3, 4, 5"
)]
pub physical_plan_type: ::core::option::Option<
ballista_physical_plan_node::PhysicalPlanType,
@@ -49,9 +49,20 @@ pub mod ballista_physical_plan_node {
UnresolvedShuffle(super::UnresolvedShuffleExecNode),
#[prost(message, tag = "4")]
SortShuffleWriter(super::SortShuffleWriterExecNode),
+ #[prost(message, tag = "5")]
+ ChaosExec(super::ChaosExecNode),
}
}
#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct ChaosExecNode {
+ #[prost(double, tag = "1")]
+ pub failure_probability: f64,
+ #[prost(string, tag = "2")]
+ pub fault_type: ::prost::alloc::string::String,
+ #[prost(uint64, tag = "3")]
+ pub seed: u64,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
pub struct ShuffleWriterExecNode {
/// TODO it seems redundant to provide job and stage id here since we also have them
/// in the TaskDefinition that wraps this plan
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index ed384625e..a231c0a87 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -53,7 +53,7 @@ use std::{convert::TryInto, io::Cursor};
use crate::execution_plans::sort_shuffle::SortShuffleConfig;
use crate::execution_plans::{
- CoalescePlan, PartitionGroup, ShuffleReaderExec, ShuffleWriterExec,
+ ChaosExec, CoalescePlan, PartitionGroup, ShuffleReaderExec, ShuffleWriterExec,
SortShuffleWriterExec, UnresolvedShuffleExec,
};
use crate::serde::protobuf::{
@@ -532,6 +532,23 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
};
Ok(Arc::new(exec))
}
+ PhysicalPlanType::ChaosExec(chaos_exec) => {
+ let input = match inputs {
+ [input] => input.clone(),
+ _ => {
+ return Err(DataFusionError::Internal(format!(
+ "ChaosExec expects exactly 1 input, got {}",
+ inputs.len()
+ )));
+ }
+ };
+ Ok(Arc::new(ChaosExec::new(
+ input,
+ chaos_exec.failure_probability,
+ &chaos_exec.fault_type,
+ Some(chaos_exec.seed),
+ )?))
+ }
}
}
@@ -690,6 +707,22 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
))
})?;
+ Ok(())
+ } else if let Some(exec) = node.as_any().downcast_ref::<ChaosExec>() {
+ let proto = protobuf::BallistaPhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::ChaosExec(
+ protobuf::ChaosExecNode {
+ failure_probability: exec.failure_probability(),
+ fault_type: exec.fault_type().to_string(),
+ seed: exec.seed(),
+ },
+ )),
+ };
+ proto.encode(buf).map_err(|e| {
+ DataFusionError::Internal(format!(
+ "failed to encode chaos monkey execution plan: {e:?}"
+ ))
+ })?;
Ok(())
} else {
Err(DataFusionError::Internal(format!(
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
new file mode 100644
index 000000000..a2b8f7b2e
--- /dev/null
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs
@@ -0,0 +1,204 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use ballista_core::config::BallistaConfig;
+use ballista_core::execution_plans::ChaosExec;
+use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion::config::ConfigOptions;
+use datafusion::error::Result;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::ExecutionPlan;
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+use std::sync::Arc;
+
+/// Optimizer rule that randomly injects a single [`ChaosExec`] into the physical plan.
+///
+/// Controlled by four Ballista config keys:
+///
+/// - `ballista.testing.chaos_execution.enabled` (bool, default `false`)
+/// - `ballista.testing.chaos_execution.probability` (f64 in 0.0–1.0, default `0.25`)
+/// - `ballista.testing.chaos_execution.fault_type` (str, default `"transient"`) — one of `transient`, `fatal`, `panic`, `delay`
+/// - `ballista.testing.chaos_execution.seed` (str, default `""`) — optional u64 seed for reproducible runs
+///
+/// When enabled, the rule walks the plan in pre-order, picks a random node, and wraps it
+/// with [`ChaosExec`]. Exactly one injection per `optimize` call.
+///
+/// This optimizer rule is not idempotent and should not be called multiple times
+/// on the same plan.
+#[derive(Debug, Default)]
+pub struct ChaosCreatingRule {}
+
+impl PhysicalOptimizerRule for ChaosCreatingRule {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let bc = config
+ .extensions
+ .get::<BallistaConfig>()
+ .cloned()
+ .unwrap_or_default();
+
+ if !bc.chaos_execution_enabled() {
+ return Ok(plan);
+ }
+
+ let failure_probability = bc.chaos_execution_probability();
+ let fault_type = bc.chaos_execution_fault_type();
+ let seed = bc.chaos_execution_seed();
+
+ // Count total nodes (pre-order DFS).
+ let mut node_count = 0usize;
+ plan.apply(|_| {
+ node_count += 1;
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+
+ // Pick a uniformly random target. % 1 == 0 always, so single-node plans
+ // are deterministic (index 0), which makes tests predictable.
+ let target_idx = match seed {
+ Some(s) => (StdRng::seed_from_u64(s).random::<u64>() as usize) % node_count,
+ None => (rand::random::<u64>() as usize) % node_count,
+ };
+ let mut current_idx = 0usize;
+
+ // transform_down visits nodes in the same pre-order as apply, so
+ // target_idx addresses the same node in both passes.
+ let result = plan.transform_down(|node| {
+ let idx = current_idx;
+ current_idx += 1;
+ if idx == target_idx {
+ let chaos =
+ Arc::new(ChaosExec::new(node, failure_probability, &fault_type, seed)?)
+ as Arc<dyn ExecutionPlan>;
+ log::warn!("A ChaosExec node has been injected in your execution plan, making execution non-deterministic");
+ Ok(Transformed::yes(chaos))
+ } else {
+ Ok(Transformed::no(node))
+ }
+ })?;
+
+ Ok(result.data)
+ }
+
+ fn name(&self) -> &str {
+ "ChaosCreatingRule"
+ }
+
+ fn schema_check(&self) -> bool {
+ false
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::assert_plan;
+ use ballista_core::config::BallistaConfig;
+ use ballista_core::execution_plans::ChaosExec;
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::common::{ColumnStatistics, Statistics};
+ use datafusion::config::{ConfigOptions, ExtensionOptions};
+ use datafusion::physical_plan::ExecutionPlan;
+ use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion::physical_plan::test::exec::StatisticsExec;
+ use std::sync::Arc;
+
+ fn leaf_exec() -> Arc<dyn ExecutionPlan> {
+ let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
+ let stats = Statistics {
+ num_rows: Default::default(),
+ total_byte_size: Default::default(),
+ column_statistics: vec![ColumnStatistics::new_unknown()],
+ };
+ Arc::new(StatisticsExec::new(stats, schema))
+ }
+
+ fn chaos_config(probability: f64, seed: u64) -> ConfigOptions {
+ let mut bc = BallistaConfig::default();
+ bc.set("testing.chaos_execution.enabled", "true").unwrap();
+ bc.set(
+ "testing.chaos_execution.probability",
+ &probability.to_string(),
+ )
+ .unwrap();
+ bc.set("testing.chaos_execution.seed", &seed.to_string())
+ .unwrap();
+ let mut config = ConfigOptions::default();
+ config.extensions.insert(bc);
+ config
+ }
+
+ #[test]
+ fn chaos_rule_disabled_is_noop() {
+ let rule = ChaosCreatingRule::default();
+ let plan = leaf_exec();
+ let result = rule.optimize(plan, &ConfigOptions::default()).unwrap();
+ assert_plan!(result.as_ref(), @ "StatisticsExec: col_count=1, row_count=Absent");
+ }
+
+ // Single-node plan: target_idx = rand % 1 = 0 always — injection is deterministic.
+ #[test]
+ fn chaos_rule_wraps_single_node() {
+ let rule = ChaosCreatingRule::default();
+ let plan = leaf_exec();
+ let result = rule.optimize(plan, &chaos_config(1.0, 442)).unwrap();
+ assert_plan!(result.as_ref(), @ r"
+ ChaosExec: failure_probability=1, fault_type=transient, seed=442
+ StatisticsExec: col_count=1, row_count=Absent
+ ");
+ }
+
+ // Multi-node plan: target node is random, so we only assert count, not shape.
+ #[test]
+ fn chaos_rule_injects_exactly_once() {
+ let rule = ChaosCreatingRule::default();
+ let plan =
+ Arc::new(CoalescePartitionsExec::new(leaf_exec())) as Arc<dyn ExecutionPlan>;
+
+ for _ in 0..20 {
+ let result = rule.optimize(plan.clone(), &chaos_config(0.5, 12)).unwrap();
+ let mut count = 0usize;
+ result
+ .apply(|node| {
+ if node.as_any().downcast_ref::<ChaosExec>().is_some() {
+ count += 1;
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })
+ .unwrap();
+ assert_eq!(1, count, "exactly one ChaosExec must be present");
+ }
+ }
+
+ #[test]
+ fn chaos_rule_probability_is_passed_through() {
+ let rule = ChaosCreatingRule::default();
+ let plan = leaf_exec();
+ let result = rule.optimize(plan, &chaos_config(0.42, 42)).unwrap();
+ let chaos = result
+ .as_any()
+ .downcast_ref::<ChaosExec>()
+ .expect("single-node plan always wraps root");
+ assert!(
+ (chaos.failure_probability() - 0.42).abs() < 1e-5,
+ "probability should round-trip through f64→f32 cast"
+ );
+ }
+}
diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
index 5838694c3..b4ae1fbc8 100644
--- a/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
+++ b/ballista/scheduler/src/state/aqe/optimizer_rule/mod.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+pub mod chaos_exec;
pub mod coalesce_partitions;
pub mod distributed_exchange;
pub mod join_selection;
diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs
index 298d77c34..ada158752 100644
--- a/ballista/scheduler/src/state/aqe/planner.rs
+++ b/ballista/scheduler/src/state/aqe/planner.rs
@@ -16,6 +16,7 @@
// under the License.
use crate::state::aqe::adapter::BallistaAdapter;
use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec};
+use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule;
use crate::state::aqe::optimizer_rule::{
CoalescePartitionsRule, DelayJoinSelectionRule, DistributedExchangeRule,
PropagateEmptyExecRule, SelectJoinRule,
@@ -506,7 +507,10 @@ impl AdaptivePlanner {
/// set of rules which will be executed ONCE before
/// running standard set of physical optimizers
fn plan_preparation_optimizers() -> Vec<PhysicalOptimizerRuleRef> {
- vec![Arc::new(DelayJoinSelectionRule::default())]
+ vec![
+ Arc::new(DelayJoinSelectionRule::default()),
+ Arc::new(ChaosCreatingRule::default()),
+ ]
}
/// Creates a session state with the given configuration and optimizer rules.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org