Check the final physical plan in the planner and name the rule that broke it

datafusion/core/src/optimizer_rule_reference.md

@@ -69,27 +69,31 @@ Rule order matters. The default pipeline may change between releases.
 The same rule name may appear more than once when the default pipeline runs it
 in multiple phases.
 
-| order | rule                           | phase                   | summary                                                                                                      |
-| ----- | ------------------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------ |
-| 1     | `OutputRequirements`           | add phase               | Adds helper nodes so output requirements survive later physical rewrites.                                    |
-| 2     | `aggregate_statistics`         | -                       | Uses exact source statistics to answer some aggregates without scanning data.                                |
-| 3     | `EliminateConstantJoin`        | -                       | Replaces inner joins against a single row of literals with a filter that keeps the other input's order.      |
-| 4     | `join_selection`               | -                       | Chooses join implementation, build side, and partition mode from statistics and stream properties.           |
-| 5     | `LimitedDistinctAggregation`   | -                       | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed.              |
-| 6     | `FilterPushdown`               | pre-optimization phase  | Pushes supported physical filters down toward data sources before distribution and sorting are enforced.     |
-| 7     | `WindowTopN`                   | -                       | Replaces eligible row-number window and filter patterns with per-partition TopK execution.                   |
-| 8     | `EnsureRequirements`           | -                       | Enforces both distribution and sorting requirements in a single idempotent rule.                             |
-| 9     | `CombinePartialFinalAggregate` | -                       | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant.             |
-| 10    | `OptimizeAggregateOrder`       | -                       | Updates aggregate expressions to use the best ordering once sort requirements are known.                     |
-| 11    | `ProjectionPushdown`           | early pass              | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure.           |
-| 12    | `OutputRequirements`           | remove phase            | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done.          |
-| 13    | `LimitAggregation`             | -                       | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets.                   |
-| 14    | `LimitPushPastWindows`         | -                       | Pushes fetch limits through bounded window operators when doing so keeps the result correct.                 |
-| 15    | `HashJoinBuffering`            | -                       | Adds buffering on the probe side of hash joins so probing can start before build completion.                 |
-| 16    | `LimitPushdown`                | -                       | Moves physical limits into child operators or fetch-enabled variants to cut data early.                      |
-| 17    | `TopKRepartition`              | -                       | Pushes TopK below hash repartition when the partition key is a prefix of the sort key.                       |
-| 18    | `ProjectionPushdown`           | late pass               | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities.               |
-| 19    | `PushdownSort`                 | -                       | Pushes sort requirements into data sources that can already return sorted output.                            |
-| 20    | `EnsureCooperative`            | -                       | Wraps non-cooperative plan parts so long-running tasks yield fairly.                                         |
-| 21    | `FilterPushdown(Post)`         | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving.                        |
-| 22    | `SanityCheckPlan`              | -                       | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. |
+After the last rule, including any rule added by the user, the physical planner
+checks that the final plan meets its ordering, distribution, and infinite-input
+safety requirements (`SanityCheckPlan`). If the ordering or distribution
+requirements do not hold, the error names the rule that broke them.
+
+| order | rule                           | phase                   | summary                                                                                                  |
+| ----- | ------------------------------ | ----------------------- | -------------------------------------------------------------------------------------------------------- |
+| 1     | `OutputRequirements`           | add phase               | Adds helper nodes so output requirements survive later physical rewrites.                                |
+| 2     | `aggregate_statistics`         | -                       | Uses exact source statistics to answer some aggregates without scanning data.                            |
+| 3     | `EliminateConstantJoin`        | -                       | Replaces inner joins against a single row of literals with a filter that keeps the other input's order.  |
+| 4     | `join_selection`               | -                       | Chooses join implementation, build side, and partition mode from statistics and stream properties.       |
+| 5     | `LimitedDistinctAggregation`   | -                       | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed.          |
+| 6     | `FilterPushdown`               | pre-optimization phase  | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. |
+| 7     | `WindowTopN`                   | -                       | Replaces eligible row-number window and filter patterns with per-partition TopK execution.               |
+| 8     | `EnsureRequirements`           | -                       | Enforces both distribution and sorting requirements in a single idempotent rule.                         |
+| 9     | `CombinePartialFinalAggregate` | -                       | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant.         |
+| 10    | `OptimizeAggregateOrder`       | -                       | Updates aggregate expressions to use the best ordering once sort requirements are known.                 |
+| 11    | `ProjectionPushdown`           | early pass              | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure.       |
+| 12    | `OutputRequirements`           | remove phase            | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done.      |
+| 13    | `LimitAggregation`             | -                       | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets.               |
+| 14    | `LimitPushPastWindows`         | -                       | Pushes fetch limits through bounded window operators when doing so keeps the result correct.             |
+| 15    | `HashJoinBuffering`            | -                       | Adds buffering on the probe side of hash joins so probing can start before build completion.             |
+| 16    | `LimitPushdown`                | -                       | Moves physical limits into child operators or fetch-enabled variants to cut data early.                  |
+| 17    | `TopKRepartition`              | -                       | Pushes TopK below hash repartition when the partition key is a prefix of the sort key.                   |
+| 18    | `ProjectionPushdown`           | late pass               | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities.           |
+| 19    | `PushdownSort`                 | -                       | Pushes sort requirements into data sources that can already return sorted output.                        |
+| 20    | `EnsureCooperative`            | -                       | Wraps non-cooperative plan parts so long-running tasks yield fairly.                                     |
+| 21    | `FilterPushdown(Post)`         | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving.                    |

datafusion/core/src/physical_planner.rs

@@ -92,6 +92,10 @@ use datafusion_physical_expr::expressions::Literal;
 use datafusion_physical_expr::{
     LexOrdering, PhysicalSortExpr, create_physical_sort_exprs,
 };
+use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements;
+use datafusion_physical_optimizer::sanity_checker::{
+    check_plan_tree_sanity, input_requirements_hold,
+};
 use datafusion_physical_plan::empty::EmptyExec;
 use datafusion_physical_plan::execution_plan::InvariantLevel;
 use datafusion_physical_plan::joins::PiecewiseMergeJoinExec;
@@ -3411,9 +3415,18 @@ impl DefaultPhysicalPlanner {
             observer(new_plan.as_ref(), optimizer.as_ref())
         }
 
-        // This runs once after all optimizer runs are complete,
-        // to verify that the plan is executable.
-        InvariantChecker(InvariantLevel::Executable).check(&new_plan)?;
+        // This runs once after all optimizer runs are complete, including
+        // rules appended by the user, to verify that the plan is executable.
+        let optimizer_options = &session_state.config_options().optimizer;
+        if let Err(e) = check_plan_tree_sanity(&new_plan, optimizer_options) {
+            return Err(sanity_check_error(
+                e,
+                &new_plan,
+                plan,
+                optimizers,
+                &optimizer_context,
+            ));
+        }
 
         debug!(
             "Optimized physical plan:\n{}\n",
@@ -3660,6 +3673,73 @@ fn tuple_err<T, R>(value: (Result<T>, Result<R>)) -> Result<(T, R)> {
     }
 }
 
+/// Adds context to an error from the final [`check_plan_tree_sanity`].
+///
+/// `EnsureRequirements` makes the order and distribution requirements hold,
+/// and every rule after it must keep them. When they do not hold for
+/// `final_plan`, the context names the rule that broke this contract: either
+/// `EnsureRequirements` itself, or the rule after which the requirements
+/// stopped holding. To find it, `optimizers` run again from `input_plan`, so
+/// planning that succeeds keeps no intermediate plans.
+fn sanity_check_error(
+    error: DataFusionError,
+    final_plan: &Arc<dyn ExecutionPlan>,
+    input_plan: Arc<dyn ExecutionPlan>,
+    optimizers: &[Arc<dyn PhysicalOptimizerRule + Send + Sync>],
+    optimizer_context: &dyn PhysicalOptimizerContext,
+) -> DataFusionError {
+    const SANITY_CHECK: &str = "SanityCheckPlan";
+    let plain_error =
+        |error| DataFusionError::Context(SANITY_CHECK.to_string(), Box::new(error));
+    if input_requirements_hold(final_plan) {
+        // The error is about unbounded inputs or another invariant.
+        return plain_error(error);
+    }
+
+    let ensure_requirements = EnsureRequirements::new();
+    let enforced = optimizers
+        .iter()
+        .rposition(|rule| rule.name() == ensure_requirements.name());
+    // Plans before `EnsureRequirements` need not satisfy the requirements, so
+    // only plans from its last run onwards are checked. Without it, every
+    // plan is checked, starting with the input plan.
+    let first_checked = enforced.unwrap_or(0);
+    let mut held = enforced.is_none() && input_requirements_hold(&input_plan);
+    let mut broken_by = None;
+    let mut plan = input_plan;
+    for (idx, rule) in optimizers.iter().enumerate() {
+        plan = match rule.optimize_with_context(plan, optimizer_context) {
+            Ok(plan) => plan,
+            // The rules did not run as they did before, so no rule is named.
+            Err(_) => return plain_error(error),
+        };
+        if idx >= first_checked {
+            let holds = input_requirements_hold(&plan);
+            if held && !holds {
+                broken_by = Some(rule);
+            }
+            held = holds;
+        }
+    }
+    let context = match (broken_by, enforced) {
+        (Some(rule), _) => format!(
+            "{SANITY_CHECK}: order and distribution requirements stopped holding \
+             after physical optimizer rule '{}'",
+            rule.name()
+        ),
+        (None, Some(_)) => format!(
+            "{SANITY_CHECK}: the plan produced by physical optimizer rule '{}' \
+             does not satisfy the order and distribution requirements",
+            ensure_requirements.name()
+        ),
+        (None, None) => format!(
+            "{SANITY_CHECK}: no physical optimizer rule produced a plan that \
+             satisfies the order and distribution requirements"
+        ),
+    };
+    DataFusionError::Context(context, Box::new(error))
+}
+
 struct OptimizationInvariantChecker<'a> {
     rule: &'a Arc<dyn PhysicalOptimizerRule + Send + Sync>,
 }

datafusion/core/tests/memory_limit/mod.rs

@@ -578,7 +578,7 @@ async fn sort_merge_join_no_spill() {
         ])
         .with_memory_limit(1_000)
         .with_config(config)
-        .with_scenario(Scenario::AccessLogStreaming)
+        .with_scenario(Scenario::AccessLogStreamingByPodTime)
         .run()
         .await
 }
@@ -597,7 +597,7 @@ async fn sort_merge_join_spill() {
         .with_memory_limit(1_000)
         .with_config(config)
         .with_disk_manager_builder(DiskManagerBuilder::default())
-        .with_scenario(Scenario::AccessLogStreaming)
+        .with_scenario(Scenario::AccessLogStreamingByPodTime)
         .with_expected_success()
         .run()
         .await
@@ -1457,6 +1457,11 @@ enum Scenario {
     /// [`StreamingTable`]
     AccessLogStreaming,
 
+    /// [`Self::AccessLogStreaming`], sorted by `pod, time`, with a
+    /// [`StreamingTable`] that declares that order. A sort merge join on
+    /// `pod, time` then needs no sort.
+    AccessLogStreamingByPodTime,
+
     /// N partitions of sorted, dictionary encoded strings.
     DictionaryStrings {
         partitions: usize,
@@ -1518,6 +1523,23 @@ impl Scenario {
                 .with_infinite_table(true);
                 Arc::new(table)
             }
+            Self::AccessLogStreamingByPodTime => {
+                let batches = access_log_batches_by_pod_time();
+                let table = StreamingTable::try_new(
+                    batches[0].schema(),
+                    vec![Arc::new(DummyStreamPartition {
+                        schema: batches[0].schema(),
+                        batches,
+                    })],
+                )
+                .unwrap()
+                .with_infinite_table(true)
+                .with_sort_order(vec![
+                    datafusion_expr::col("pod").sort(true, true),
+                    datafusion_expr::col("time").sort(true, true),
+                ]);
+                Arc::new(table)
+            }
             Self::DictionaryStrings {
                 partitions,
                 single_row_batches,
@@ -1566,7 +1588,7 @@ impl Scenario {
                 // has a memory budget which we'll likely hit first)
                 Some(vec![Arc::new(JoinSelection::new())])
             }
-            Self::AccessLogStreaming => {
+            Self::AccessLogStreaming | Self::AccessLogStreamingByPodTime => {
                 // Disable all physical optimizer rules except the
                 // JoinSelection rule to avoid sorts or repartition,
                 // as they also have memory budgets that may be hit
@@ -1642,6 +1664,26 @@ fn access_log_batches() -> Vec<RecordBatch> {
         .collect()
 }
 
+/// [`access_log_batches`] sorted by `pod, time`, in batches of 50 rows
+fn access_log_batches_by_pod_time() -> Vec<RecordBatch> {
+    let batches = access_log_batches();
+    let batch = arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap();
+    let sort_column = |name: &str| arrow::compute::SortColumn {
+        values: Arc::clone(batch.column_by_name(name).unwrap()),
+        options: Some(SortOptions::default()),
+    };
+    let indices = arrow::compute::lexsort_to_indices(
+        &[sort_column("pod"), sort_column("time")],
+        None,
+    )
+    .unwrap();
+    let batch = arrow::compute::take_record_batch(&batch, &indices).unwrap();
+    (0..batch.num_rows())
+        .step_by(50)
+        .map(|offset| batch.slice(offset, 50.min(batch.num_rows() - offset)))
+        .collect()
+}
+
 /// If `one_row_batches` is true, then returns new record batches that
 /// are one row in size
 fn maybe_split_batches(