Add Expr::try_as_col, deprecate Expr::try_into_col (#10448)

Author: alamb

datafusion/expr/src/expr.rs

@@ -1264,13 +1264,36 @@ impl Expr {
         })
     }
 
+    #[deprecated(since = "39.0.0", note = "use try_as_col instead")]
     pub fn try_into_col(&self) -> Result<Column> {
         match self {
             Expr::Column(it) => Ok(it.clone()),
             _ => plan_err!("Could not coerce '{self}' into Column!"),
         }
     }
 
+    /// 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<HashSet<Column>> {
         let mut using_columns = HashSet::new();

datafusion/expr/src/expr_rewriter/mod.rs

@@ -221,7 +221,7 @@ pub fn coerce_plan_expr_for_schema(
             let exprs: Vec<Expr> = 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))

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::<HashSet<Expr>>();
         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::<Result<Vec<_>>>()?;
 

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);

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::<Vec<_>>();

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();

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::<Result<Vec<_>, _>>()?;
                         builder.join_using(
                             into_logical_plan!(join.right, ctx, extension_codec)?,