From 86cfd8956a0182c7f3c867e7db2cf23fffcaf22c Mon Sep 17 00:00:00 2001
From: Xin Li <bluer906@gmail.com>
Date: Tue, 23 Jul 2024 21:24:11 +0800
Subject: [PATCH] Fix comment

---
 datafusion/expr/src/expr.rs              | 20 ++++++++++++++++++
 datafusion/expr/src/logical_plan/plan.rs | 26 ++++++++++++------------
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
index e3620501d9a8..452c05be34f4 100644
--- a/datafusion/expr/src/expr.rs
+++ b/datafusion/expr/src/expr.rs
@@ -1340,6 +1340,9 @@ impl Expr {
     ///
     /// returns `None` if the expression is not a `Column`
     ///
+    /// Note: None may be returned for expressions that are not `Column` but
+    /// are convertible to `Column` such as `Cast` expressions.
+    ///
     /// Example
     /// ```
     /// # use datafusion_common::Column;
@@ -1358,6 +1361,23 @@ impl Expr {
         }
     }
 
+    /// Returns the inner `Column` if any. This is a specialized version of
+    /// [`Self::try_as_col`] that take Cast expressions into account when the
+    /// expression is as on condition for joins.
+    ///
+    /// Called this method when you are sure that the expression is a `Column`
+    /// or a `Cast` expression that wraps a `Column`.
+    pub fn get_as_join_column(&self) -> Option<&Column> {
+        match self {
+            Expr::Column(c) => Some(c),
+            Expr::Cast(Cast { expr, .. }) => match &**expr {
+                Expr::Column(c) => Some(c),
+                _ => None,
+            },
+            _ => None,
+        }
+    }
+
     /// Return all referenced columns of this expression.
     #[deprecated(since = "40.0.0", note = "use Expr::column_refs instead")]
     pub fn to_columns(&self) -> Result<HashSet<Column>> {
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
index 9c22a9d6d07f..d4fe233cac06 100644
--- a/datafusion/expr/src/logical_plan/plan.rs
+++ b/datafusion/expr/src/logical_plan/plan.rs
@@ -35,7 +35,7 @@ use crate::utils::{
     grouping_set_expr_count, grouping_set_to_exprlist, split_conjunction,
 };
 use crate::{
-    build_join_schema, expr_vec_fmt, BinaryExpr, BuiltInWindowFunction, Cast,
+    build_join_schema, expr_vec_fmt, BinaryExpr, BuiltInWindowFunction,
     CreateMemoryTable, CreateView, Expr, ExprSchemable, LogicalPlanBuilder, Operator,
     TableProviderFilterPushDown, TableSource, WindowFunctionDefinition,
 };
@@ -496,10 +496,18 @@ impl LogicalPlan {
                 // The join keys in using-join must be columns.
                 let columns =
                     on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| {
-                        let l = Self::as_col(l.clone())?;
-                        let r = Self::as_col(r.clone())?;
-                        accumu.insert(l);
-                        accumu.insert(r);
+                        let Some(l) = l.get_as_join_column() else {
+                            return internal_err!(
+                                "Invalid join key. Expected column, found {l:?}"
+                            );
+                        };
+                        let Some(r) = r.get_as_join_column() else {
+                            return internal_err!(
+                                "Invalid join key. Expected column, found {r:?}"
+                            );
+                        };
+                        accumu.insert(l.to_owned());
+                        accumu.insert(r.to_owned());
                         Result::<_, DataFusionError>::Ok(accumu)
                     })?;
                 using_columns.push(columns);
@@ -510,14 +518,6 @@ impl LogicalPlan {
         Ok(using_columns)
     }
 
-    fn as_col(expr: Expr) -> Result<Column> {
-        match expr {
-            Expr::Column(c) => Ok(c),
-            Expr::Cast(Cast { expr, .. }) => Self::as_col(*expr),
-            _ => internal_err!("Invalid join key. Expected column, found {expr:?}"),
-        }
-    }
-
     /// returns the first output expression of this `LogicalPlan` node.
     pub fn head_output_expr(&self) -> Result<Option<Expr>> {
         match self {