Start cleaning up tests

alamb · Mar 31, 2024 · a569474 · a569474
1 parent ebb7073
commit a569474
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 12 deletions.
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -1297,7 +1297,8 @@ mod test {
         let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
         dbg!(&plan);
         let expected =
-            "Projection: CAST(Utf8(\"1998-03-18\") AS Timestamp(Nanosecond, None)) = CAST(CAST(Utf8(\"1998-03-18\") AS Date32) AS Timestamp(Nanosecond, None))\n  EmptyRelation";
+            "Projection: CAST(Utf8(\"1998-03-18\") AS Timestamp(Nanosecond, None)) = CAST(CAST(Utf8(\"1998-03-18\") AS Date32) AS Timestamp(Nanosecond, None))\
+            \n  EmptyRelation";
         assert_analyzed_plan_eq(Arc::new(TypeCoercion::new()), &plan, expected)?;
         Ok(())
     }

diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs
@@ -535,12 +535,12 @@ mod tests {
         assert_eq!(
             "Optimizer rule 'get table_scan rule' failed\ncaused by\nget table_scan rule\ncaused by\n\
             Internal error: Failed due to a difference in schemas, \
-            original schema: DFSchema { fields: [\
+            original schema: DFSchema { fields: [], metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }, \
+            new schema: DFSchema { fields: [\
             DFField { qualifier: Some(Bare { table: \"test\" }), field: Field { name: \"a\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, \
             DFField { qualifier: Some(Bare { table: \"test\" }), field: Field { name: \"b\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, \
             DFField { qualifier: Some(Bare { table: \"test\" }), field: Field { name: \"c\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }], \
-            metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }, \
-            new schema: DFSchema { fields: [], metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }.\
+            metadata: {}, functional_dependencies: FunctionalDependencies { deps: [] } }.\
             \nThis was likely caused by a bug in DataFusion's code and we would welcome that you file an bug report in our issue tracker",
             err.strip_backtrace()
         );

diff --git a/datafusion/optimizer/src/push_down_projection.rs b/datafusion/optimizer/src/push_down_projection.rs
@@ -502,7 +502,8 @@ mod tests {
             .build()?;
 
         let optimized_plan1 = optimize(plan).expect("failed to optimize plan");
-        let optimized_plan2 = optimize(optimized_plan1).expect("failed to optimize plan");
+        let optimized_plan2 =
+            optimize(optimized_plan1.clone()).expect("failed to optimize plan");
 
         let formatted_plan1 = format!("{optimized_plan1:?}");
         let formatted_plan2 = format!("{optimized_plan2:?}");

diff --git a/datafusion/optimizer/src/test/mod.rs b/datafusion/optimizer/src/test/mod.rs
@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::analyzer::{Analyzer, AnalyzerRule};
-use crate::optimizer::{assert_only_schema_is_the_same, Optimizer};
+use crate::optimizer::Optimizer;
 use crate::{OptimizerContext, OptimizerRule};
 use arrow::datatypes::{DataType, Field, Schema};
 use datafusion_common::config::ConfigOptions;
@@ -156,13 +156,9 @@ pub fn assert_optimized_plan_eq(
     expected: &str,
 ) -> Result<()> {
     fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
-    let starting_schema = plan.schema();
 
     let optimizer = Optimizer::with_rules(vec![rule.clone()]);
     let optimized_plan = optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
-
-    // Ensure schemas always match after an optimization
-    assert_only_schema_is_the_same(rule.name(), &starting_schema, &optimized_plan)?;
     let formatted_plan = format!("{optimized_plan:?}");
     assert_eq!(formatted_plan, expected);
 
@@ -184,7 +180,6 @@ pub fn assert_optimized_plan_eq_with_rules(
         .expect("failed to optimize plan");
     let formatted_plan = format!("{optimized_plan:?}");
     assert_eq!(formatted_plan, expected);
-    assert_eq!(plan.schema(), optimized_plan.schema());
     Ok(())
 }
 
@@ -239,7 +234,7 @@ pub fn assert_optimization_skipped(
     plan: LogicalPlan,
 ) -> Result<()> {
     let optimizer = Optimizer::with_rules(vec![rule]);
-    let new_plan = optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
+    let new_plan = optimizer.optimize(plan.clone(), &OptimizerContext::new(), observe)?;
 
     assert_eq!(
         format!("{}", plan.display_indent()),