diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index f0f41a4c55c5..660a45c27a29 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -1264,6 +1264,7 @@ impl Expr { }) } + #[deprecated(since = "39.0.0", note = "use try_as_col instead")] pub fn try_into_col(&self) -> Result { match self { Expr::Column(it) => Ok(it.clone()), @@ -1271,6 +1272,28 @@ impl Expr { } } + /// Return a reference to the inner `Column` if any + /// + /// returns `None` if the expression is not a `Column` + /// + /// Example + /// ``` + /// # use datafusion_common::Column; + /// use datafusion_expr::{col, Expr}; + /// let expr = col("foo"); + /// assert_eq!(expr.try_as_col(), Some(&Column::from("foo"))); + /// + /// let expr = col("foo").alias("bar"); + /// assert_eq!(expr.try_as_col(), None); + /// ``` + pub fn try_as_col(&self) -> Option<&Column> { + if let Expr::Column(it) = self { + Some(it) + } else { + None + } + } + /// Return all referenced columns of this expression. pub fn to_columns(&self) -> Result> { let mut using_columns = HashSet::new(); diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index 700dd560ec0b..1441374bdba3 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -221,7 +221,7 @@ pub fn coerce_plan_expr_for_schema( let exprs: Vec = plan.schema().iter().map(Expr::from).collect(); let new_exprs = coerce_exprs_for_schema(exprs, plan.schema(), schema)?; - let add_project = new_exprs.iter().any(|expr| expr.try_into_col().is_err()); + let add_project = new_exprs.iter().any(|expr| expr.try_as_col().is_none()); if add_project { let projection = Projection::try_new(new_exprs, Arc::new(plan.clone()))?; Ok(LogicalPlan::Projection(projection)) diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 3f15b84784f1..2c6cfd8f9d20 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -1489,7 +1489,7 @@ pub fn wrap_projection_for_join_if_necessary( let mut projection = expand_wildcard(input_schema, &input, None)?; let join_key_items = alias_join_keys .iter() - .flat_map(|expr| expr.try_into_col().is_err().then_some(expr)) + .flat_map(|expr| expr.try_as_col().is_none().then_some(expr)) .cloned() .collect::>(); projection.extend(join_key_items); @@ -1504,8 +1504,12 @@ pub fn wrap_projection_for_join_if_necessary( let join_on = alias_join_keys .into_iter() .map(|key| { - key.try_into_col() - .or_else(|_| Ok(Column::from_name(key.display_name()?))) + if let Some(col) = key.try_as_col() { + Ok(col.clone()) + } else { + let name = key.display_name()?; + Ok(Column::from_name(name)) + } }) .collect::>>()?; diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index 9832b69f841a..266e7abc341a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -369,8 +369,18 @@ impl LogicalPlan { // The join keys in using-join must be columns. let columns = on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| { - accumu.insert(l.try_into_col()?); - accumu.insert(r.try_into_col()?); + let Some(l) = l.try_as_col().cloned() else { + return internal_err!( + "Invalid join key. Expected column, found {l:?}" + ); + }; + let Some(r) = r.try_as_col().cloned() else { + return internal_err!( + "Invalid join key. Expected column, found {r:?}" + ); + }; + accumu.insert(l); + accumu.insert(r); Result::<_, DataFusionError>::Ok(accumu) })?; using_columns.push(columns); diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 9ce135b0d646..57b38bd0d0fd 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -535,8 +535,8 @@ fn push_down_join( .on .iter() .filter_map(|(l, r)| { - let left_col = l.try_into_col().ok()?; - let right_col = r.try_into_col().ok()?; + let left_col = l.try_as_col().cloned()?; + let right_col = r.try_as_col().cloned()?; Some((left_col, right_col)) }) .collect::>(); diff --git a/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs index 9dcb8ed15563..c8638eb72395 100644 --- a/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs @@ -52,7 +52,7 @@ impl TreeNodeRewriter for ShortenInListSimplifier { // expressions list.len() == 1 || list.len() <= THRESHOLD_INLINE_INLIST - && expr.try_into_col().is_ok() + && expr.try_as_col().is_some() ) { let first_val = list[0].clone(); diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index a6352bcefc3e..83e58c3a22cc 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -45,8 +45,9 @@ use datafusion::{ prelude::SessionContext, }; use datafusion_common::{ - context, internal_err, not_impl_err, parsers::CompressionTypeVariant, - plan_datafusion_err, DataFusionError, Result, TableReference, + context, internal_datafusion_err, internal_err, not_impl_err, + parsers::CompressionTypeVariant, plan_datafusion_err, DataFusionError, Result, + TableReference, }; use datafusion_expr::{ dml, @@ -695,7 +696,12 @@ impl AsLogicalPlan for LogicalPlanNode { // The equijoin keys in using-join must be column. let using_keys = left_keys .into_iter() - .map(|key| key.try_into_col()) + .map(|key| { + key.try_as_col().cloned() + .ok_or_else(|| internal_datafusion_err!( + "Using join keys must be column references, got: {key:?}" + )) + }) .collect::, _>>()?; builder.join_using( into_logical_plan!(join.right, ctx, extension_codec)?,