Skip to content

Commit

Permalink
Minor: reduce cloneing in infer_placeholder_types (#5638)
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb authored Mar 20, 2023
1 parent d77ccc2 commit 89f2102
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 34 deletions.
10 changes: 6 additions & 4 deletions datafusion/common/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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!()))
};
}

Expand Down Expand Up @@ -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)]
Expand Down
57 changes: 27 additions & 30 deletions datafusion/sql/src/expr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down Expand Up @@ -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)
})
Expand Down

0 comments on commit 89f2102

Please sign in to comment.