[arrow-datafusion] branch main updated: Minor: reduce cloneing in `infer_placeholder_types` (#5638)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 89f210266 Minor: reduce cloneing in `infer_placeholder_types` (#5638)
89f210266 is described below

commit 89f2102666a2dd235147386856dd76e2079c2d97
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Mon Mar 20 12:25:40 2023 +0100

    Minor: reduce cloneing in `infer_placeholder_types` (#5638)
---
 datafusion/common/src/error.rs | 10 +++++---
 datafusion/sql/src/expr/mod.rs | 57 ++++++++++++++++++++----------------------
 2 files changed, 33 insertions(+), 34 deletions(-)

diff --git a/datafusion/common/src/error.rs b/datafusion/common/src/error.rs
index 0231895f7..c1e11fe93 100644
--- a/datafusion/common/src/error.rs
+++ b/datafusion/common/src/error.rs
@@ -98,10 +98,7 @@ pub enum DataFusionError {
 #[macro_export]
 macro_rules! context {
     ($desc:expr, $err:expr) => {
-        datafusion_common::DataFusionError::Context(
-            format!("{} at {}:{}", $desc, file!(), line!()),
-            Box::new($err),
-        )
+        $err.context(format!("{} at {}:{}", $desc, file!(), line!()))
     };
 }
 
@@ -417,6 +414,11 @@ impl DataFusionError {
         // return last checkpoint (which may be the original error)
         last_datafusion_error
     }
+
+    /// wraps self in Self::Context with a description
+    pub fn context(self, description: impl Into<String>) -> Self {
+        Self::Context(description.into(), Box::new(self))
+    }
 }
 
 #[cfg(test)]
diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs
index 555344166..daad118e3 100644
--- a/datafusion/sql/src/expr/mod.rs
+++ b/datafusion/sql/src/expr/mod.rs
@@ -78,7 +78,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         let mut expr = self.sql_expr_to_logical_expr(sql, schema, planner_context)?;
         expr = self.rewrite_partial_qualifier(expr, schema);
         self.validate_schema_satisfies_exprs(schema, &[expr.clone()])?;
-        let expr = infer_placeholder_types(expr, schema.clone())?;
+        let expr = infer_placeholder_types(expr, schema)?;
         Ok(expr)
     }
 
@@ -499,36 +499,33 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
     }
 }
 
-/// Find all `PlaceHolder` tokens in a logical plan, and try to infer their type from context
-fn infer_placeholder_types(expr: Expr, schema: DFSchema) -> Result<Expr> {
-    rewrite_expr(expr, |expr| {
-        let expr = match expr {
-            Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
-                let left = (*left).clone();
-                let right = (*right).clone();
-                let lt = left.get_type(&schema);
-                let rt = right.get_type(&schema);
-                let left = match (&left, rt) {
-                    (Expr::Placeholder { id, data_type }, Ok(dt)) => Expr::Placeholder {
-                        id: id.clone(),
-                        data_type: Some(data_type.clone().unwrap_or(dt)),
-                    },
-                    _ => left.clone(),
-                };
-                let right = match (&right, lt) {
-                    (Expr::Placeholder { id, data_type }, Ok(dt)) => Expr::Placeholder {
-                        id: id.clone(),
-                        data_type: Some(data_type.clone().unwrap_or(dt)),
-                    },
-                    _ => right.clone(),
-                };
-                Expr::BinaryExpr(BinaryExpr {
-                    left: Box::new(left),
-                    op,
-                    right: Box::new(right),
-                })
+// modifies expr if it is a placeholder with datatype of right
+fn rewrite_placeholder(expr: &mut Expr, other: &Expr, schema: &DFSchema) -> Result<()> {
+    if let Expr::Placeholder { id: _, data_type } = expr {
+        if data_type.is_none() {
+            let other_dt = other.get_type(schema);
+            match other_dt {
+                Err(e) => {
+                    return Err(e.context(format!(
+                        "Can not find type of {other} needed to infer type of {expr}"
+                    )))?;
+                }
+                Ok(dt) => {
+                    *data_type = Some(dt);
+                }
             }
-            _ => expr.clone(),
+        };
+    }
+    Ok(())
+}
+
+/// Find all [`Expr::PlaceHolder`] tokens in a logical plan, and try to infer their type from context
+fn infer_placeholder_types(expr: Expr, schema: &DFSchema) -> Result<Expr> {
+    rewrite_expr(expr, |mut expr| {
+        // Default to assuming the arguments are the same type
+        if let Expr::BinaryExpr(BinaryExpr { left, op: _, right }) = &mut expr {
+            rewrite_placeholder(left.as_mut(), right.as_ref(), schema)?;
+            rewrite_placeholder(right.as_mut(), left.as_ref(), schema)?;
         };
         Ok(expr)
     })