Replace panic in datafusion-expr crate (#3341)

apache · Sep 8, 2022 · e6378f4 · e6378f4
1 parent 30fce22
commit e6378f4
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 25 deletions.
diff --git a/datafusion/expr/src/binary_rule.rs b/datafusion/expr/src/binary_rule.rs
@@ -382,10 +382,10 @@ fn coercion_decimal_mathematics_type(
                     let result_precision = result_scale + (*p1 - *s1).min(*p2 - *s2);
                     Some(create_decimal_type(result_precision, result_scale))
                 }
-                _ => unreachable!(),
+                _ => None,
             }
         }
-        _ => unreachable!(),
+        _ => None,
     }
 }
 

diff --git a/datafusion/expr/src/expr_rewriter.rs b/datafusion/expr/src/expr_rewriter.rs
@@ -386,7 +386,7 @@ fn rewrite_sort_col_by_aggs(expr: Expr, plan: &LogicalPlan) -> Result<Expr> {
                         // The expr is not based on Aggregate plan output. Skip it.
                         return Ok(expr);
                     }
-                    let normalized_expr = normalized_expr.unwrap();
+                    let normalized_expr = normalized_expr?;
                     if let Some(found_agg) = self
                         .aggr_expr
                         .iter()

diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs
@@ -607,12 +607,12 @@ impl LogicalPlanBuilder {
         for (l, r) in &on {
             if self.plan.schema().field_from_column(l).is_ok()
                 && right.schema().field_from_column(r).is_ok()
-                && can_hash(self.plan.schema().field_from_column(l).unwrap().data_type())
+                && can_hash(self.plan.schema().field_from_column(l)?.data_type())
             {
                 join_on.push((l.clone(), r.clone()));
             } else if self.plan.schema().field_from_column(r).is_ok()
                 && right.schema().field_from_column(l).is_ok()
-                && can_hash(self.plan.schema().field_from_column(r).unwrap().data_type())
+                && can_hash(self.plan.schema().field_from_column(r)?.data_type())
             {
                 join_on.push((r.clone(), l.clone()));
             } else {
@@ -629,7 +629,9 @@ impl LogicalPlanBuilder {
         }
         if join_on.is_empty() {
             let join = Self::from(self.plan.clone()).cross_join(&right.clone())?;
-            join.filter(filters.unwrap())
+            join.filter(filters.ok_or_else(|| {
+                DataFusionError::Internal("filters should not be None here".to_string())
+            })?)
         } else {
             Ok(Self::from(LogicalPlan::Join(Join {
                 left: Arc::new(self.plan.clone()),
@@ -901,18 +903,16 @@ pub fn union_with_alias(
         .map(|p| match p.as_ref() {
             LogicalPlan::Projection(Projection {
                 expr, input, alias, ..
-            }) => Arc::new(
-                project_with_column_index_alias(
-                    expr.to_vec(),
-                    input.clone(),
-                    union_schema.clone(),
-                    alias.clone(),
-                )
-                .unwrap(),
-            ),
-            x => Arc::new(x.clone()),
+            }) => Ok(Arc::new(project_with_column_index_alias(
+                expr.to_vec(),
+                input.clone(),
+                union_schema.clone(),
+                alias.clone(),
+            )?)),
+            x => Ok(Arc::new(x.clone())),
         })
-        .collect::<Vec<_>>();
+        .into_iter()
+        .collect::<Result<Vec<_>>>()?;
 
     if inputs.is_empty() {
         return Err(DataFusionError::Plan("Empty UNION".to_string()));

diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs
@@ -231,8 +231,12 @@ impl<'a, 'b> PlanVisitor for GraphvizVisitor<'a, 'b> {
         _plan: &LogicalPlan,
     ) -> std::result::Result<bool, fmt::Error> {
         // always be non-empty as pre_visit always pushes
-        self.parent_ids.pop().unwrap();
-        Ok(true)
+        // So it should always be Ok(true)
+        let res = self.parent_ids.pop();
+        match res {
+            Some(_) => Ok(true),
+            None => Err(fmt::Error),
+        }
     }
 }
 

diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
@@ -549,8 +549,10 @@ impl LogicalPlan {
             fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                 let with_schema = false;
                 let mut visitor = IndentVisitor::new(f, with_schema);
-                self.0.accept(&mut visitor).unwrap();
-                Ok(())
+                match self.0.accept(&mut visitor) {
+                    Ok(_) => Ok(()),
+                    Err(_) => Err(fmt::Error),
+                }
             }
         }
         Wrapper(self)
@@ -590,8 +592,10 @@ impl LogicalPlan {
             fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                 let with_schema = true;
                 let mut visitor = IndentVisitor::new(f, with_schema);
-                self.0.accept(&mut visitor).unwrap();
-                Ok(())
+                match self.0.accept(&mut visitor) {
+                    Ok(_) => Ok(()),
+                    Err(_) => Err(fmt::Error),
+                }
             }
         }
         Wrapper(self)
@@ -641,12 +645,12 @@ impl LogicalPlan {
                 let mut visitor = GraphvizVisitor::new(f);
 
                 visitor.pre_visit_plan("LogicalPlan")?;
-                self.0.accept(&mut visitor).unwrap();
+                self.0.accept(&mut visitor).map_err(|_| fmt::Error)?;
                 visitor.post_visit_plan()?;
 
                 visitor.set_with_schema(true);
                 visitor.pre_visit_plan("Detailed LogicalPlan")?;
-                self.0.accept(&mut visitor).unwrap();
+                self.0.accept(&mut visitor).map_err(|_| fmt::Error)?;
                 visitor.post_visit_plan()?;
 
                 writeln!(f, "}}")?;