Fix Internal Error for an INNER JOIN query

xinlifoobar · Jul 21, 2024 · 39bf7db · 39bf7db
1 parent 1b9b35c
commit 39bf7db
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 13 deletions.
diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs
@@ -35,9 +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,
-    CreateMemoryTable, CreateView, Expr, ExprSchemable, LogicalPlanBuilder, Operator,
-    TableProviderFilterPushDown, TableSource, WindowFunctionDefinition,
+    build_join_schema, expr_vec_fmt, BinaryExpr, BuiltInWindowFunction, Cast, CreateMemoryTable, CreateView, Expr, ExprSchemable, LogicalPlanBuilder, Operator, TableProviderFilterPushDown, TableSource, WindowFunctionDefinition
 };
 
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
@@ -496,16 +494,8 @@ impl LogicalPlan {
                 // The join keys in using-join must be columns.
                 let columns =
                     on.iter().try_fold(HashSet::new(), |mut accumu, (l, r)| {
-                        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:?}"
-                            );
-                        };
+                        let l = Self::as_col(l.clone())?;
+                        let r = Self::as_col(r.clone())?;
                         accumu.insert(l);
                         accumu.insert(r);
                         Result::<_, DataFusionError>::Ok(accumu)
@@ -518,6 +508,14 @@ 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 {

diff --git a/datafusion/sqllogictest/test_files/join.slt b/datafusion/sqllogictest/test_files/join.slt
@@ -1054,3 +1054,45 @@ DROP TABLE t1;
 
 statement ok
 DROP TABLE t2;
+
+# Join Using Issue with Cast Expr
+# Found issue: https://github.com/apache/datafusion/issues/11412
+
+statement ok
+/*DML*/CREATE TABLE t60(v0 BIGINT, v1 BIGINT, v2 BOOLEAN, v3 BOOLEAN);
+
+statement ok
+/*DML*/CREATE TABLE t0(v0 DOUBLE, v1 BIGINT);
+
+statement ok
+/*DML*/CREATE TABLE t1(v0 DOUBLE);
+
+query I
+SELECT COUNT(*)
+FROM t1
+NATURAL JOIN t60
+INNER JOIN t0
+ON t60.v1 = t0.v0
+AND t0.v1 > t60.v1;
+----
+0
+
+query I
+SELECT COUNT(*)
+FROM t1
+JOIN t60
+USING (v0)
+INNER JOIN t0
+ON t60.v1 = t0.v0
+AND t0.v1 > t60.v1;
+----
+0
+
+statement ok
+DROP TABLE t60;
+
+statement ok
+DROP TABLE t0;
+
+statement ok
+DROP TABLE t1;