alamb commented on code in PR #26461:
URL: https://github.com/apache/datafusion/pull/26461#discussion_r4417178213


##########
datafusion/core/src/physical_planner.rs:
##########
@@ -3391,6 +3395,10 @@ impl DefaultPhysicalPlanner {
 
         let mut new_plan = Arc::clone(&plan);
         let optimizer_context = SessionOptimizerContext::new(session_state);
+        // The plan after each rule, so that a failed final check can name the
+        // rule that broke the plan. Plans share unchanged subtrees, so this
+        // costs one pointer per rule.

Review Comment:
   it also pins the plan memory in place for longer, so it would make peak 
memory usage higher which maybe doesn't matter, but it is a differene



##########
datafusion/core/tests/physical_optimizer/sanity_checker.rs:
##########
@@ -999,3 +1003,116 @@ async fn test_union_with_sorts_and_constants() -> 
Result<()> {
 
     Ok(())
 }
+
+/// Removes every [`SortExec`], standing in for a rule that breaks ordering

Review Comment:
   FWIW I was confused at first why this is needed. It is a test it turns out. 
Maybe calling it `MockRule` or adding a comment like '/// Test rule to check 
messages' would be more clear



##########
docs/source/library-user-guide/upgrading/56.0.0.md:
##########
@@ -1336,3 +1336,34 @@ without declaring one, use 
`FileScanConfigBuilder::with_preserve_order(true)`.
 Code that reads `FileScanConfig::preserve_order` to decide whether files may be
 reordered should call `FileScanConfig::preserves_file_order()` instead, which
 also accounts for the output ordering.
+
+### `SanityCheckPlan` is deprecated and no longer in the default physical 
optimizer rules
+
+The physical planner now checks the final physical plan itself, after the last

Review Comment:
   I think we could probably shorten this and note that the `SanityCheckPlan` 
has been replaced with a method and if custom pipelines need to check the plan 
they should call the new method directly



##########
datafusion/core/tests/physical_optimizer/sanity_checker.rs:
##########
@@ -999,3 +1003,116 @@ async fn test_union_with_sorts_and_constants() -> 
Result<()> {
 
     Ok(())
 }
+
+/// Removes every [`SortExec`], standing in for a rule that breaks ordering
+/// requirements.
+///
+/// [`SortExec`]: datafusion_physical_plan::sorts::sort::SortExec
+#[derive(Debug)]
+struct RemoveSorts;
+
+impl PhysicalOptimizerRule for RemoveSorts {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        plan.transform_up(|plan| {
+            if plan.is::<SortExec>() {
+                Ok(Transformed::yes(Arc::clone(plan.children()[0])))
+            } else {
+                Ok(Transformed::no(plan))
+            }
+        })
+        .data()
+    }
+
+    fn name(&self) -> &str {
+        "remove_sorts"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
+async fn plan_window_query(state: SessionState) -> Result<Arc<dyn 
ExecutionPlan>> {
+    let ctx = SessionContext::new_with_state(state);
+    ctx.sql("CREATE TABLE t(a INT) AS VALUES (3), (1), (2)")
+        .await?
+        .collect()
+        .await?;
+    ctx.sql("SELECT a, row_number() OVER (ORDER BY a) AS rn FROM t")
+        .await?
+        .create_physical_plan()
+        .await
+}
+
+/// A rule appended after the default rules is checked, and the error names it.
+#[tokio::test]
+async fn test_final_check_names_appended_rule() -> Result<()> {
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_physical_optimizer_rule(Arc::new(RemoveSorts))
+        .build();
+    let error = plan_window_query(state).await.unwrap_err().to_string();
+    assert!(
+        error.starts_with(
+            "SanityCheckPlan: order and distribution requirements stopped 
holding \
+             after physical optimizer rule 'remove_sorts'"
+        ),
+        "{error}"
+    );
+    assert!(
+        error.contains("does not satisfy order requirements"),
+        "{error}"
+    );
+    Ok(())
+}
+
+/// A custom rule list is checked too.
+#[tokio::test]
+async fn test_final_check_names_rule_in_custom_pipeline() -> Result<()> {

Review Comment:
   How is this different than test_final_check_names_appended_rule ?



##########
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)

Review Comment:
   This change seems unrelated to the rest of this PR? Or on more review it 
seems related to forcing an error. However it is really strange to me that we 
are reusing the tests for memory limitiing to test optimizer errors. I would 
suggest we move the optimizer error stuff into an optimizer_integration style 
test



##########
datafusion/core/src/physical_planner.rs:
##########
@@ -3344,12 +3352,16 @@ impl DefaultPhysicalPlanner {
                 optimizer.name(),
                 displayable(new_plan.as_ref()).indent(false)
             );
-            observer(new_plan.as_ref(), optimizer.as_ref())
+            observer(new_plan.as_ref(), optimizer.as_ref());
+            plan_after_rule.push((optimizer, Arc::clone(&new_plan)));
         }
 
-        // 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) {

Review Comment:
   I actually think we shouldn't allow plans to opt out of plan sanity checking 
-- if the generated plan doesn't pass sanity check then it will fail in some 
other way (either wrong results or worse). Catching it early is better I think



##########
datafusion/core/src/physical_planner.rs:
##########
@@ -3391,6 +3395,10 @@ impl DefaultPhysicalPlanner {
 
         let mut new_plan = Arc::clone(&plan);
         let optimizer_context = SessionOptimizerContext::new(session_state);
+        // The plan after each rule, so that a failed final check can name the
+        // rule that broke the plan. Plans share unchanged subtrees, so this
+        // costs one pointer per rule.

Review Comment:
   Also there are a bunch of rules that try to `unwrap` the Arc to avoid 
allocations which this will prevent.
   
   If you only want to get the plans on failure, I recommend saving only the 
input plan and if there is a failure then go back and run each of the passes 
(again) to get a more specific message



##########
datafusion/sqllogictest/test_files/explain.slt:
##########
@@ -242,7 +242,6 @@ physical_plan after ProjectionPushdown SAME TEXT AS ABOVE
 physical_plan after PushdownSort SAME TEXT AS ABOVE
 physical_plan after EnsureCooperative SAME TEXT AS ABOVE
 physical_plan after FilterPushdown(Post) SAME TEXT AS ABOVE
-physical_plan after SanityCheckPlan SAME TEXT AS ABOVE

Review Comment:
   I do like that this is no longer an OptimizerRule



##########
datafusion/core/tests/physical_optimizer/sanity_checker.rs:
##########
@@ -999,3 +1003,116 @@ async fn test_union_with_sorts_and_constants() -> 
Result<()> {
 
     Ok(())
 }
+
+/// Removes every [`SortExec`], standing in for a rule that breaks ordering
+/// requirements.
+///
+/// [`SortExec`]: datafusion_physical_plan::sorts::sort::SortExec
+#[derive(Debug)]
+struct RemoveSorts;
+
+impl PhysicalOptimizerRule for RemoveSorts {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        plan.transform_up(|plan| {
+            if plan.is::<SortExec>() {
+                Ok(Transformed::yes(Arc::clone(plan.children()[0])))
+            } else {
+                Ok(Transformed::no(plan))
+            }
+        })
+        .data()
+    }
+
+    fn name(&self) -> &str {
+        "remove_sorts"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}
+
+async fn plan_window_query(state: SessionState) -> Result<Arc<dyn 
ExecutionPlan>> {
+    let ctx = SessionContext::new_with_state(state);
+    ctx.sql("CREATE TABLE t(a INT) AS VALUES (3), (1), (2)")
+        .await?
+        .collect()
+        .await?;
+    ctx.sql("SELECT a, row_number() OVER (ORDER BY a) AS rn FROM t")
+        .await?
+        .create_physical_plan()
+        .await
+}
+
+/// A rule appended after the default rules is checked, and the error names it.
+#[tokio::test]
+async fn test_final_check_names_appended_rule() -> Result<()> {
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_physical_optimizer_rule(Arc::new(RemoveSorts))
+        .build();
+    let error = plan_window_query(state).await.unwrap_err().to_string();
+    assert!(
+        error.starts_with(
+            "SanityCheckPlan: order and distribution requirements stopped 
holding \
+             after physical optimizer rule 'remove_sorts'"
+        ),
+        "{error}"
+    );
+    assert!(
+        error.contains("does not satisfy order requirements"),
+        "{error}"
+    );
+    Ok(())
+}
+
+/// A custom rule list is checked too.
+#[tokio::test]
+async fn test_final_check_names_rule_in_custom_pipeline() -> Result<()> {
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_physical_optimizer_rules(vec![
+            Arc::new(EnsureRequirements::new()),
+            Arc::new(RemoveSorts),
+        ])
+        .build();
+    let error = plan_window_query(state).await.unwrap_err().to_string();
+    assert!(
+        error.starts_with(
+            "SanityCheckPlan: order and distribution requirements stopped 
holding \
+             after physical optimizer rule 'remove_sorts'"
+        ),
+        "{error}"
+    );
+    Ok(())
+}
+
+/// Without `EnsureRequirements`, the window's sort is never added.
+#[tokio::test]
+async fn test_final_check_without_ensure_requirements() -> Result<()> {
+    let state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_physical_optimizer_rules(vec![Arc::new(RemoveSorts)])
+        .build();
+    let error = plan_window_query(state).await.unwrap_err().to_string();
+    assert!(
+        error.starts_with(
+            "SanityCheckPlan: no physical optimizer rule produced a plan that \
+             satisfies the order and distribution requirements"
+        ),
+        "{error}"
+    );
+    Ok(())
+}
+
+/// The default rules produce a plan that passes the final check.
+#[tokio::test]
+async fn test_final_check_passes_default_rules() -> Result<()> {

Review Comment:
   this I think is redundant as it is tested many times in the rest of the code



-- 
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]

Reply via email to