Improve join type support in substrait

datafusion/substrait/src/consumer.rs

@@ -38,7 +38,7 @@ use substrait::proto::{
     },
     extensions::simple_extension_declaration::MappingType,
     function_argument::ArgType,
-    plan_rel,
+    join_rel, plan_rel,
     read_rel::ReadType,
     rel::RelType,
     sort_field::{SortDirection, SortKind::*},
@@ -317,19 +317,7 @@ pub async fn from_substrait_rel(
             let right = LogicalPlanBuilder::from(
                 from_substrait_rel(ctx, join.right.as_ref().unwrap(), extensions).await?,
             );
-            let join_type = match join.r#type {
-                1 => JoinType::Inner,
-                2 => JoinType::Left,
-                3 => JoinType::Right,
-                4 => JoinType::Full,
-                5 => JoinType::LeftAnti,
-                6 => JoinType::LeftSemi,
-                _ => {
-                    return Err(DataFusionError::Internal(
-                        "invalid join type".to_string(),
-                    ))
-                }
-            };
+            let join_type = from_substrait_jointype(join.r#type)?;
             let schema =
                 build_join_schema(left.schema(), right.schema(), &JoinType::Inner)?;
             let on = from_substrait_rex(
@@ -432,6 +420,30 @@ pub async fn from_substrait_rel(
     }
 }
 
+fn from_substrait_jointype(join_type: i32) -> Result<JoinType> {
+    if let Some(substrait_join_type) = join_rel::JoinType::from_i32(join_type) {
+        match substrait_join_type {
+            join_rel::JoinType::Inner => Ok(JoinType::Inner),
+            join_rel::JoinType::Left => Ok(JoinType::Left),
+            join_rel::JoinType::Right => Ok(JoinType::Right),
+            join_rel::JoinType::Outer => Ok(JoinType::Full),
+            join_rel::JoinType::Anti => Ok(JoinType::LeftAnti),
+            join_rel::JoinType::Semi => Ok(JoinType::LeftSemi),
+            _ => {
+                return Err(DataFusionError::Internal(format!(
+                    "unsupported join type {:?}",
+                    substrait_join_type
+                )))
+            }
+        }
+    } else {
+        return Err(DataFusionError::Internal(format!(
+            "invalid join type variant {:?}",
+            join_type
+        )));
+    }
+}
+
 /// Convert Substrait AggregateFunction to DataFusion Expr
 pub async fn from_substrait_agg_func(
     f: &AggregateFunction,

datafusion/substrait/src/producer.rs

@@ -45,7 +45,7 @@ use substrait::proto::{
         simple_extension_declaration::{ExtensionFunction, MappingType},
     },
     function_argument::ArgType,
-    plan_rel,
+    join_rel, plan_rel,
     read_rel::{NamedTable, ReadType},
     rel::RelType,
     sort_field::{SortDirection, SortKind},
@@ -243,15 +243,7 @@ pub fn to_substrait_rel(
         LogicalPlan::Join(join) => {
             let left = to_substrait_rel(join.left.as_ref(), extension_info)?;
             let right = to_substrait_rel(join.right.as_ref(), extension_info)?;
-            let join_type = match join.join_type {
-                JoinType::Inner => 1,
-                JoinType::Left => 2,
-                JoinType::Right => 3,
-                JoinType::Full => 4,
-                JoinType::LeftAnti => 5,
-                JoinType::LeftSemi => 6,
-                _ => panic!(), // TODO
-            };
+            let join_type = to_substrait_jointype(join.join_type);
             // we only support basic joins so return an error for anything not yet supported
             if join.null_equals_null {
                 return Err(DataFusionError::NotImplemented(
@@ -282,7 +274,7 @@ pub fn to_substrait_rel(
                         common: None,
                         left: Some(left),
                         right: Some(right),
-                        r#type: join_type,
+                        r#type: join_type as i32,
                         expression: Some(Box::new(to_substrait_rex(
                             &e,
                             &join.schema,
@@ -310,6 +302,18 @@ pub fn to_substrait_rel(
     }
 }
 
+fn to_substrait_jointype(join_type: JoinType) -> join_rel::JoinType {
+    match join_type {
+        JoinType::Inner => join_rel::JoinType::Inner,
+        JoinType::Left => join_rel::JoinType::Left,
+        JoinType::Right => join_rel::JoinType::Right,
+        JoinType::Full => join_rel::JoinType::Outer,
+        JoinType::LeftAnti => join_rel::JoinType::Anti,
+        JoinType::LeftSemi => join_rel::JoinType::Semi,
+        JoinType::RightAnti | JoinType::RightSemi => unimplemented!(),
+    }
+}
+
 pub fn operator_to_name(op: Operator) -> &'static str {
     match op {
         Operator::Eq => "equal",

datafusion/substrait/tests/roundtrip.rs

@@ -182,6 +182,22 @@ mod tests {
         .await
     }
 
+    #[tokio::test]
+    async fn roundtrip_left_join() -> Result<()> {
+        roundtrip("SELECT data.a FROM data LEFT JOIN data2 ON data.a = data2.a").await
+    }
+
+    #[tokio::test]
+    async fn roundtrip_right_join() -> Result<()> {
+        roundtrip("SELECT data.a FROM data RIGHT JOIN data2 ON data.a = data2.a").await
+    }
+
+    #[tokio::test]
+    async fn roundtrip_outer_join() -> Result<()> {
+        roundtrip("SELECT data.a FROM data FULL OUTER JOIN data2 ON data.a = data2.a")
+            .await
+    }
+
     async fn assert_expected_plan(sql: &str, expected_plan_str: &str) -> Result<()> {
         let mut ctx = create_context().await?;
         let df = ctx.sql(sql).await?;