comphead commented on code in PR #2801: URL: https://github.com/apache/datafusion-ballista/pull/2801#discussion_r4575372173
########## ballista/scheduler/src/state/aqe/plan_change_log.rs: ########## @@ -0,0 +1,415 @@ +// 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. + +//! Logs how the adaptive planner's rules change a job's plan. +//! +//! With `debug` logging on for this module, each pass of rules over a job's +//! plan logs every rule that changed the plan, with a diff of the plan before +//! and after the rule, and then the plan the pass ended with. A rule that +//! leaves the plan as it was logs nothing, as with LLVM's `-print-changed` or +//! Spark's plan change log, and neither does a change a later rule of the +//! pass undoes, such as the node DataFusion's `OutputRequirements` adds and +//! removes again: +//! +//! ```text +//! RUST_LOG=ballista_scheduler::state::aqe::plan_change_log=debug +//! ``` +//! +//! What a rule decided, and why it skipped an optimization, the rule logs +//! under its own module, so `RUST_LOG=ballista_scheduler::state::aqe=debug` +//! shows both. +//! +//! Some rules change an exchange in place, so plans are compared by their +//! rendering, which is only taken while this log is enabled. A pass logs its +//! changes once its rules have run, or have failed. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{DefaultHasher, Hash, Hasher}; + +use ballista_core::JobId; +use datafusion::physical_plan::{ExecutionPlan, displayable}; +use log::{Level, debug, log_enabled}; + +/// Unchanged lines a diff keeps on each side of a change. +const CONTEXT: usize = 3; + +/// The largest table of line matches a diff builds, `(m + 1) * (n + 1)` for +/// the `m` and `n` lines between the first and last change of each plan, about +/// 500 lines each. Past it, all of those lines show as removed and added. +const MAX_MATCH_TABLE: usize = 1 << 18; + +/// What the planner runs a pass of rules over a job's plan for. +#[derive(Clone, Copy, Debug)] +pub(crate) enum Pass { + /// The first optimization of the job's plan. + Initial, + /// The replan after a stage finished. + StageFinished(usize), + /// The replan after a stage about to launch read a finished stage instead. + StageReused, + /// The rules run on a stage's own plan as it launches. + Stage(usize), +} + +impl Display for Pass { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Initial => write!(f, "initial optimization"), + Self::StageFinished(stage_id) => { + write!(f, "replan after stage {stage_id} finished") + } + Self::StageReused => write!(f, "replan after reusing a finished stage"), + Self::Stage(stage_id) => write!(f, "stage {stage_id}"), + } + } +} + +/// Logs the rules of one [`Pass`] that changed a job's plan. See the module +/// docs. +pub(crate) struct PlanChangeLog<'a> { + job_id: &'a JobId, + pass: Pass, + /// What the rules of the pass changed so far, or `None` while the log is + /// disabled. + changes: Option<Changes>, +} + +/// The changes the rules of a pass made to a plan. +struct Changes { + /// The plan as the last rule left it, rendered. + plan: String, + /// Each change no later rule undid: the rule's name and its diff. + made: Vec<(String, String)>, + /// A digest of each plan the pass went through, `seen[k]` being the plan + /// the first `k` changes of `made` led to, so a rule that brings a plan + /// back drops the changes it undid. + seen: Vec<u64>, +} + +impl<'a> PlanChangeLog<'a> { + /// Starts the log of `pass` over `plan`. + pub(crate) fn new(job_id: &'a JobId, pass: Pass, plan: &dyn ExecutionPlan) -> Self { + Self::with_enabled(log_enabled!(Level::Debug), job_id, pass, plan) + } + + fn with_enabled( + enabled: bool, + job_id: &'a JobId, + pass: Pass, + plan: &dyn ExecutionPlan, + ) -> Self { + let changes = enabled.then(|| { + let plan = render(plan); + Changes { + seen: vec![digest(&plan)], + plan, + made: Vec::new(), + } + }); + Self { + job_id, + pass, + changes, + } + } + + /// Records the change `rule` made, given the `plan` it returned, if it + /// changed the plan. + pub(crate) fn record(&mut self, rule: &str, plan: &dyn ExecutionPlan) { + let Some(changes) = &mut self.changes else { + return; + }; + let after = render(plan); + let after_digest = digest(&after); + match changes.seen.iter().position(|seen| *seen == after_digest) { + // The rule undid every change since the plan was last like this, + // or changed nothing. + Some(back) => { + changes.made.truncate(back); + changes.seen.truncate(back + 1); + } + None => { + changes + .made + .push((rule.to_owned(), diff(&changes.plan, &after))); + changes.seen.push(after_digest); + } + } + changes.plan = after; + } +} + +impl Drop for PlanChangeLog<'_> { + /// Logs the changes of the pass, and the plan they left, if any. A pass + /// that fails logs what its rules changed before it failed. + fn drop(&mut self) { + let Some(changes) = &self.changes else { + return; + }; + if changes.made.is_empty() { + return; + } + for (rule, diff) in &changes.made { + debug!( + "Job {}, {}: {rule} changed the plan:\n{diff}", + self.job_id, self.pass + ); + } + let rules: Vec<&str> = + changes.made.iter().map(|(rule, _)| rule.as_str()).collect(); + debug!( + "Job {}, {}: plan changed by {}:\n{}", + self.job_id, + self.pass, + rules.join(", "), + changes.plan.trim_end() + ); + } +} + +fn render(plan: &dyn ExecutionPlan) -> String { + displayable(plan).indent(false).to_string() +} + +fn digest(rendering: &str) -> u64 { + let mut hasher = DefaultHasher::new(); + rendering.hash(&mut hasher); + hasher.finish() +} + +/// A line diff of two renderings of a plan: `-` marks a removed line and `+` +/// an added one, and up to [`CONTEXT`] unchanged lines are kept on each side +/// of a change, with `...` for those left out. Lines match ignoring their +/// indentation, so a subtree a rule moves under a new node shows as the one +/// added line. +fn diff(before: &str, after: &str) -> String { + let before: Vec<&str> = before.lines().collect(); + let after: Vec<&str> = after.lines().collect(); + let mut lines = edit_script(&before, &after, str::trim_start); + if lines.iter().all(|(tag, _)| *tag == ' ') { + // Only the indentation changed: a node moved to another parent without + // moving in the rendering. + lines = edit_script(&before, &after, |line| line); + } + + let mut kept = vec![false; lines.len()]; + for (i, (tag, _)) in lines.iter().enumerate() { + if *tag != ' ' { + let end = (i + CONTEXT + 1).min(lines.len()); + kept[i.saturating_sub(CONTEXT)..end].fill(true); + } + } + let mut out = Vec::new(); + let mut skipped = false; + for ((tag, line), kept) in lines.into_iter().zip(kept) { + if !kept { + skipped = true; + } else { + if std::mem::take(&mut skipped) { + out.push("...".to_owned()); + } + out.push(format!("{tag}{line}")); + } + } + if skipped { + out.push("...".to_owned()); + } + out.join("\n") +} + +/// The lines of `before` and `after` in order, each tagged ` ` when kept, as +/// it is in `after`, `-` when removed and `+` when added. Two lines match when +/// their `key`s are equal. +fn edit_script<'a>( Review Comment: Reuse before adding: `similar` 2.7.0 is already in `Cargo.lock` (through `insta`). `similar::capture_diff_slices` over the trimmed lines, mapping each op back to the original line by index, plus its context-grouping helpers, would replace `edit_script`, the `MAX_MATCH_TABLE` fallback and most of `diff`. It also uses Myers, O(ND), instead of an O(mn) table that allocates up to 1 MiB per change. The trade-off is promoting an already-vetted transitive dev dependency to a normal dependency of the scheduler. That seems cheaper than maintaining about 100 lines of diff code and its five tests. ########## ballista/scheduler/src/state/aqe/planner.rs: ########## @@ -498,15 +486,25 @@ impl AdaptivePlanner { let config = self.session_state.config().options(); let (stage_ids, shuffle_writers) = stages .into_iter() - .map(|plan| { + .map(|mut plan| { + let mut changes = PlanChangeLog::new( + &self.job_id, + Pass::Stage(launched_stage_id(&plan)), Review Comment: `launched_stage_id(&plan)` runs on every stage launch even with the log disabled, and it adds a new `expect` panic path to the scheduler for a debug-only label. The invariant holds today because `identify_runnable_stages` assigns ids first. Still, `Pass::Stage(Option<usize>)`, rendering a missing id as `?`, would keep logging from ever taking the scheduler down. Computing it only inside the enabled branch would also work. ########## ballista/scheduler/src/state/aqe/plan_change_log.rs: ########## @@ -0,0 +1,415 @@ +// 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. + +//! Logs how the adaptive planner's rules change a job's plan. +//! +//! With `debug` logging on for this module, each pass of rules over a job's +//! plan logs every rule that changed the plan, with a diff of the plan before +//! and after the rule, and then the plan the pass ended with. A rule that +//! leaves the plan as it was logs nothing, as with LLVM's `-print-changed` or +//! Spark's plan change log, and neither does a change a later rule of the +//! pass undoes, such as the node DataFusion's `OutputRequirements` adds and +//! removes again: +//! +//! ```text +//! RUST_LOG=ballista_scheduler::state::aqe::plan_change_log=debug +//! ``` +//! +//! What a rule decided, and why it skipped an optimization, the rule logs +//! under its own module, so `RUST_LOG=ballista_scheduler::state::aqe=debug` +//! shows both. +//! +//! Some rules change an exchange in place, so plans are compared by their +//! rendering, which is only taken while this log is enabled. A pass logs its +//! changes once its rules have run, or have failed. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{DefaultHasher, Hash, Hasher}; + +use ballista_core::JobId; +use datafusion::physical_plan::{ExecutionPlan, displayable}; +use log::{Level, debug, log_enabled}; + +/// Unchanged lines a diff keeps on each side of a change. +const CONTEXT: usize = 3; + +/// The largest table of line matches a diff builds, `(m + 1) * (n + 1)` for +/// the `m` and `n` lines between the first and last change of each plan, about +/// 500 lines each. Past it, all of those lines show as removed and added. +const MAX_MATCH_TABLE: usize = 1 << 18; + +/// What the planner runs a pass of rules over a job's plan for. +#[derive(Clone, Copy, Debug)] +pub(crate) enum Pass { + /// The first optimization of the job's plan. + Initial, + /// The replan after a stage finished. + StageFinished(usize), + /// The replan after a stage about to launch read a finished stage instead. + StageReused, + /// The rules run on a stage's own plan as it launches. + Stage(usize), +} + +impl Display for Pass { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::Initial => write!(f, "initial optimization"), + Self::StageFinished(stage_id) => { + write!(f, "replan after stage {stage_id} finished") + } + Self::StageReused => write!(f, "replan after reusing a finished stage"), + Self::Stage(stage_id) => write!(f, "stage {stage_id}"), + } + } +} + +/// Logs the rules of one [`Pass`] that changed a job's plan. See the module +/// docs. +pub(crate) struct PlanChangeLog<'a> { + job_id: &'a JobId, + pass: Pass, + /// What the rules of the pass changed so far, or `None` while the log is + /// disabled. + changes: Option<Changes>, +} + +/// The changes the rules of a pass made to a plan. +struct Changes { + /// The plan as the last rule left it, rendered. + plan: String, + /// Each change no later rule undid: the rule's name and its diff. + made: Vec<(String, String)>, + /// A digest of each plan the pass went through, `seen[k]` being the plan + /// the first `k` changes of `made` led to, so a rule that brings a plan + /// back drops the changes it undid. + seen: Vec<u64>, +} + +impl<'a> PlanChangeLog<'a> { + /// Starts the log of `pass` over `plan`. + pub(crate) fn new(job_id: &'a JobId, pass: Pass, plan: &dyn ExecutionPlan) -> Self { + Self::with_enabled(log_enabled!(Level::Debug), job_id, pass, plan) + } + + fn with_enabled( + enabled: bool, + job_id: &'a JobId, + pass: Pass, + plan: &dyn ExecutionPlan, + ) -> Self { + let changes = enabled.then(|| { + let plan = render(plan); + Changes { + seen: vec![digest(&plan)], + plan, + made: Vec::new(), + } + }); + Self { + job_id, + pass, + changes, + } + } + + /// Records the change `rule` made, given the `plan` it returned, if it + /// changed the plan. + pub(crate) fn record(&mut self, rule: &str, plan: &dyn ExecutionPlan) { + let Some(changes) = &mut self.changes else { + return; + }; + let after = render(plan); + let after_digest = digest(&after); + match changes.seen.iter().position(|seen| *seen == after_digest) { + // The rule undid every change since the plan was last like this, + // or changed nothing. + Some(back) => { + changes.made.truncate(back); Review Comment: Devil's advocate on dropping undone changes: this also hides one rule overriding another's decision. `changes_a_later_rule_undoes_are_dropped` shows it. `B` changes the plan and `C` puts it back, and neither is logged. For AQE that's exactly the interaction a debugger wants to see, for example a join strategy picked by one rule and reverted by a later one. Prior art logs every effective rule. Spark's `PlanChangeLogger.logRule` logs whenever `!newPlan.fastEquals(oldPlan)` and never suppresses later reverts. It offers `spark.sql.planChangeLog.rules` to filter instead. LLVM's `-print-changed` and DataFusion's `EXPLAIN VERBOSE` behave the same way. The noise here comes from one known pair, `OutputRequirements` adding and removing its node. Two narrower options: collapse only a change undone by the very next recorded change, or keep the suppression but log a one-line `A's change was undone by B` so the interaction stays visible. ########## ballista/scheduler/src/state/aqe/planner.rs: ########## @@ -443,36 +443,24 @@ impl AdaptivePlanner { self.finalise_stage(stage_id, partitions, Some(column_statistics)) } - /// Replans the stages by applying physical optimizations. + /// Replans the stages by applying physical optimizations, logging each + /// rule that changes the plan (see [`PlanChangeLog`]). /// /// # Returns /// A `Result` indicating success or failure. - fn replan_stages(&mut self) -> common::Result<()> { + fn replan_stages(&mut self, pass: Pass) -> common::Result<()> { let interval = Instant::now(); - let mut plan = self.plan.clone(); - - debug!( - "Distributed physical plan (before optimization):\n{}\n", - displayable(plan.as_ref()).indent(false) - ); - - plan = self.physical_planner.optimize_physical_plan( - plan.clone(), + let mut changes = PlanChangeLog::new(&self.job_id, pass, self.plan.as_ref()); + let plan = self.physical_planner.optimize_physical_plan( + Arc::clone(&self.plan), &self.session_state, - |_, _| {}, + |plan, rule| changes.record(rule.name(), plan), )?; - - debug!( - "Distributed physical plan (after optimization):\n{}\n", - displayable(plan.as_ref()).indent(false) - ); - + // Log the changes before the timing below. + drop(changes); self.plan = plan; - debug!( - "Total time taken for stage replan: {:?}", - interval.elapsed() - ); + debug!("Job {}: {pass} took {:?}", self.job_id, interval.elapsed()); Review Comment: When the plan change log is on, this timing includes rendering the plan after every rule and building the diffs, because `record` runs inside the observer. Replan times get inflated exactly when someone is debugging. A note in the tuning guide would cover it. Alternatively, take `elapsed()` before `drop(changes)` and have the log measure its own overhead separately. ########## ballista/scheduler/src/state/aqe/test/reuse_exchange.rs: ########## @@ -119,16 +124,27 @@ fn aggregates_the_table(stage: &Arc<dyn ExecutionPlan>) -> bool { #[tokio::test] async fn a_volatile_subplan_runs_for_each_reader() { - // `random()` below the shuffle, where reading one stage would give both - // readers the same values. - let sql = "with v as (select c, sum(a * random()) as s from t group by c) \ - select c, s from v union all select c, s from v"; - let mut planner = planner(&context(true), sql).await; + let mut planner = planner(&context(true), VOLATILE).await; let stages = planner.identify_runnable_stages().unwrap().unwrap(); assert_eq!(stages.len(), 2); } +#[tokio::test] +async fn a_volatile_subplan_is_unshareable_for_its_volatile_expression() { Review Comment: This builds the same planner from the same `VOLATILE` query as `a_volatile_subplan_runs_for_each_reader` just above. Asserting the `Unshareable::Volatile` reason inside that test would cover both without a second planner setup. ########## ballista/scheduler/src/state/aqe/optimizer_rule/optimize_skewed_join.rs: ########## @@ -215,9 +272,10 @@ fn build_ranges( config.skew_join_skewed_partition_factor(), config.skew_join_partition_threshold_bytes() as u64, ); + let largest = sizes.iter().copied().max().unwrap_or(0); Review Comment: Nit: this is another pass over `sizes`, and `skew_threshold` just sorted a copy of them. `largest` is only used in the skip message. Returning the max from `skew_threshold` (the last element of its sorted copy), or tracking it in the `zip(sizes).map(..)` below, keeps `build_ranges` single-pass. ########## ballista/scheduler/src/state/aqe/reuse_exchange.rs: ########## @@ -77,20 +78,49 @@ use crate::state::aqe::execution_plan::ExchangeExec; /// What a stage is matched on: a digest of its encoding. See the module docs. pub(crate) type ReuseKey = [u8; 32]; +/// Why the stage of an exchange may not be shared. See the module docs. +#[derive(Debug)] +pub(crate) enum Unshareable { + /// The exchange range-partitions its output. + RangePartitioned, + /// The named operator of the stage has a volatile expression. + Volatile(String), + /// The named operator of the stage needs the stage to have one consumer. + SingleConsumer(String), + /// The stage could not be keyed, such as when the codec can't encode it. + NoKey(DataFusionError), +} + +impl Display for Unshareable { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Self::RangePartitioned => write!(f, "it range-partitions its output"), + Self::Volatile(name) => write!(f, "{name} has a volatile expression"), + Self::SingleConsumer(name) => { + write!(f, "{name} needs its stage to have one consumer") + } + Self::NoKey(e) => write!(f, "its stage could not be keyed: {e}"), + } + } +} + /// The key `node`, an exchange about to launch its stage, is matched on, or -/// `None` when its stage may not be shared. See the module docs. +/// why its stage may not be shared. See the module docs. Panics if `node` is +/// not an [`ExchangeExec`]. pub(crate) fn reuse_key( node: &Arc<dyn ExecutionPlan>, codec: &Arc<dyn PhysicalExtensionCodec>, -) -> Option<ReuseKey> { - let exchange = node.downcast_ref::<ExchangeExec>()?; +) -> std::result::Result<ReuseKey, Unshareable> { + let exchange = node + .downcast_ref::<ExchangeExec>() + .expect("only exchanges are keyed"); Review Comment: This turns a non-exchange input from `None` into a panic. The only caller (`planner.rs`, in the reuse loop) has already downcast `exchange`, so taking `exchange: &ExchangeExec` next to `node` would make the precondition a type instead of an `expect`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
