From 53f2479c337497ba4b677c9b8cf551fe0c6bcfca Mon Sep 17 00:00:00 2001
From: Andrew Lamb <andrew@nerdnetworks.org>
Date: Mon, 23 Sep 2024 12:03:22 -0400
Subject: [PATCH] Improve SanityChecker error message

---
 .../src/physical_optimizer/sanity_checker.rs  |  7 +++--
 .../physical-expr-common/src/sort_expr.rs     | 28 +++++++++++++++++--
 datafusion/physical-expr/src/partitioning.rs  | 19 +++++++++++--
 3 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/datafusion/core/src/physical_optimizer/sanity_checker.rs b/datafusion/core/src/physical_optimizer/sanity_checker.rs
index e392105fbcb7..11af51b53b72 100644
--- a/datafusion/core/src/physical_optimizer/sanity_checker.rs
+++ b/datafusion/core/src/physical_optimizer/sanity_checker.rs
@@ -33,6 +33,7 @@ use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supp
 use datafusion_physical_plan::joins::SymmetricHashJoinExec;
 use datafusion_physical_plan::{get_plan_string, ExecutionPlanProperties};
 
+use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use itertools::izip;
 
@@ -130,9 +131,9 @@ pub fn check_plan_sanity(
             if !child_eq_props.ordering_satisfy_requirement(sort_req) {
                 let plan_str = get_plan_string(&plan);
                 return plan_err!(
-                    "Plan: {:?} does not satisfy order requirements: {:?}. Child-{} order: {:?}",
+                    "Plan: {:?} does not satisfy order requirements: {}. Child-{} order: {}",
                     plan_str,
-                    sort_req,
+                    format_physical_sort_requirement_list(sort_req),
                     idx,
                     child_eq_props.oeq_class
                 );
@@ -145,7 +146,7 @@ pub fn check_plan_sanity(
         {
             let plan_str = get_plan_string(&plan);
             return plan_err!(
-                "Plan: {:?} does not satisfy distribution requirements: {:?}. Child-{} output partitioning: {:?}",
+                "Plan: {:?} does not satisfy distribution requirements: {}. Child-{} output partitioning: {}",
                 plan_str,
                 dist_req,
                 idx,
diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs
index be183b72b38b..704cb291335f 100644
--- a/datafusion/physical-expr-common/src/sort_expr.rs
+++ b/datafusion/physical-expr-common/src/sort_expr.rs
@@ -17,7 +17,7 @@
 
 //! Sort expressions
 
-use std::fmt::Display;
+use std::fmt::{Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::ops::Deref;
 use std::sync::Arc;
@@ -252,13 +252,37 @@ impl PartialEq for PhysicalSortRequirement {
     }
 }
 
-impl std::fmt::Display for PhysicalSortRequirement {
+impl Display for PhysicalSortRequirement {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         let opts_string = self.options.as_ref().map_or("NA", to_str);
         write!(f, "{} {}", self.expr, opts_string)
     }
 }
 
+/// Writes a list of [`PhysicalSortRequirement`]s to a `std::fmt::Formatter`.
+///
+/// Example output: `[a + 1, b]`
+pub fn format_physical_sort_requirement_list(
+    exprs: &[PhysicalSortRequirement],
+) -> impl Display + '_ {
+    struct DisplayWrapper<'a>(&'a [PhysicalSortRequirement]);
+    impl<'a> Display for DisplayWrapper<'a> {
+        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+            let mut iter = self.0.iter();
+            write!(f, "[")?;
+            if let Some(expr) = iter.next() {
+                write!(f, "{}", expr)?;
+            }
+            for expr in iter {
+                write!(f, ", {}", expr)?;
+            }
+            write!(f, "]")?;
+            Ok(())
+        }
+    }
+    DisplayWrapper(exprs)
+}
+
 impl PhysicalSortRequirement {
     /// Creates a new requirement.
     ///
diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs
index 272585f89815..01f72a8efd9a 100644
--- a/datafusion/physical-expr/src/partitioning.rs
+++ b/datafusion/physical-expr/src/partitioning.rs
@@ -17,13 +17,14 @@
 
 //! [`Partitioning`] and [`Distribution`] for `ExecutionPlans`
 
-use std::fmt;
-use std::sync::Arc;
-
 use crate::{
     equivalence::ProjectionMapping, expressions::UnKnownColumn, physical_exprs_equal,
     EquivalenceProperties, PhysicalExpr,
 };
+use datafusion_physical_expr_common::physical_expr::format_physical_expr_list;
+use std::fmt;
+use std::fmt::Display;
+use std::sync::Arc;
 
 /// Output partitioning supported by [`ExecutionPlan`]s.
 ///
@@ -264,6 +265,18 @@ impl Distribution {
     }
 }
 
+impl Display for Distribution {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Distribution::UnspecifiedDistribution => write!(f, "Unspecified"),
+            Distribution::SinglePartition => write!(f, "SinglePartition"),
+            Distribution::HashPartitioned(exprs) => {
+                write!(f, "HashPartitioned[{}])", format_physical_expr_list(exprs))
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {