rm unwrap cast

apache · Mar 5, 2025 · 5322d53 · 5322d53
1 parent 915abf1
commit 5322d53
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 1,439 deletions.
diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs
@@ -355,7 +355,8 @@ async fn csv_explain_verbose() {
 async fn csv_explain_inlist_verbose() {
     let ctx = SessionContext::new();
     register_aggregate_csv_by_sql(&ctx).await;
-    let sql = "EXPLAIN VERBOSE SELECT c1 FROM aggregate_test_100 where c2 in (1,2,4)";
+    // Inlist len <=3 case will be transformed to OR List so we test with len=4
+    let sql = "EXPLAIN VERBOSE SELECT c1 FROM aggregate_test_100 where c2 in (1,2,4,5)";
     let actual = execute(&ctx, sql).await;
 
     // Optimized by PreCastLitInComparisonExpressions rule
@@ -368,12 +369,12 @@ async fn csv_explain_inlist_verbose() {
     // before optimization (Int64 literals)
     assert_contains!(
         &actual,
-        "aggregate_test_100.c2 IN ([Int64(1), Int64(2), Int64(4)])"
+        "aggregate_test_100.c2 IN ([Int64(1), Int64(2), Int64(4), Int64(5)])"
     );
     // after optimization (casted to Int8)
     assert_contains!(
         &actual,
-        "aggregate_test_100.c2 IN ([Int8(1), Int8(2), Int8(4)])"
+        "aggregate_test_100.c2 IN ([Int8(1), Int8(2), Int8(4), Int8(5)])"
     );
 }
 

diff --git a/datafusion/optimizer/src/lib.rs b/datafusion/optimizer/src/lib.rs
@@ -60,7 +60,6 @@ pub mod replace_distinct_aggregate;
 pub mod scalar_subquery_to_join;
 pub mod simplify_expressions;
 pub mod single_distinct_to_groupby;
-pub mod unwrap_cast_in_comparison;
 pub mod utils;
 
 #[cfg(test)]

diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs
@@ -54,7 +54,6 @@ use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
 use crate::scalar_subquery_to_join::ScalarSubqueryToJoin;
 use crate::simplify_expressions::SimplifyExpressions;
 use crate::single_distinct_to_groupby::SingleDistinctToGroupBy;
-use crate::unwrap_cast_in_comparison::UnwrapCastInComparison;
 use crate::utils::log_plan;
 
 /// `OptimizerRule`s transforms one [`LogicalPlan`] into another which
@@ -243,7 +242,6 @@ impl Optimizer {
         let rules: Vec<Arc<dyn OptimizerRule + Sync + Send>> = vec![
             Arc::new(EliminateNestedUnion::new()),
             Arc::new(SimplifyExpressions::new()),
-            Arc::new(UnwrapCastInComparison::new()),
             Arc::new(ReplaceDistinctWithAggregate::new()),
             Arc::new(EliminateJoin::new()),
             Arc::new(DecorrelatePredicateSubquery::new()),
@@ -266,7 +264,6 @@ impl Optimizer {
             // The previous optimizations added expressions and projections,
             // that might benefit from the following rules
             Arc::new(SimplifyExpressions::new()),
-            Arc::new(UnwrapCastInComparison::new()),
             Arc::new(CommonSubexprEliminate::new()),
             Arc::new(EliminateGroupByConstant::new()),
             Arc::new(OptimizeProjections::new()),

diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
@@ -1774,7 +1774,7 @@ impl<S: SimplifyInfo> TreeNodeRewriter for Simplifier<'_, S> {
                     return internal_err!("Expect cast expr, but got {:?}", left)?;
                 };
 
-                let expr_type = info.get_data_type(&left_expr)?;
+                let expr_type = info.get_data_type(left_expr)?;
                 let right_exprs = list
                     .into_iter()
                     .map(|right| {
@@ -1851,19 +1851,17 @@ fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary<S: SimplifyInfo
             }),
             Expr::Literal(lit_val),
         ) => {
-            let Ok(expr_type) = info.get_data_type(&left_expr) else {
+            let Ok(expr_type) = info.get_data_type(left_expr) else {
                 return false;
             };
 
             let Ok(lit_type) = info.get_data_type(literal) else {
                 return false;
             };
 
-            if let Some(_) = try_cast_literal_to_type(lit_val, &expr_type) {
-                return is_supported_type(&expr_type) && is_supported_type(&lit_type);
-            }
-
-            false
+            try_cast_literal_to_type(lit_val, &expr_type).is_some()
+                && is_supported_type(&expr_type)
+                && is_supported_type(&lit_type)
         }
         _ => false,
     }
@@ -1884,7 +1882,7 @@ fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist<S: SimplifyInfo
         return false;
     };
 
-    let Ok(expr_type) = info.get_data_type(&left_expr) else {
+    let Ok(expr_type) = info.get_data_type(left_expr) else {
         return false;
     };
 
@@ -1902,12 +1900,8 @@ fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist<S: SimplifyInfo
         }
 
         match right {
-            Expr::Literal(lit_val) => {
-                if let Some(_) = try_cast_literal_to_type(lit_val, &expr_type) {
-                    continue;
-                }
-                return false;
-            }
+            Expr::Literal(lit_val)
+                if try_cast_literal_to_type(lit_val, &expr_type).is_some() => {}
             _ => return false,
         }
     }