substrait support

Signed-off-by: Ruihang Xia <[email protected]>

Author: waynexia

datafusion/substrait/src/logical_plan/consumer/expr/subquery.rs

@@ -16,11 +16,12 @@
 // under the License.
 
 use crate::logical_plan::consumer::SubstraitConsumer;
-use datafusion::common::{substrait_err, DFSchema, Spans};
-use datafusion::logical_expr::expr::{Exists, InSubquery};
-use datafusion::logical_expr::{Expr, Subquery};
+use datafusion::common::{substrait_datafusion_err, substrait_err, DFSchema, Spans};
+use datafusion::logical_expr::expr::{Exists, InSubquery, SetComparison, SetQuantifier};
+use datafusion::logical_expr::{Expr, Operator, Subquery};
 use std::sync::Arc;
 use substrait::proto::expression as substrait_expression;
+use substrait::proto::expression::subquery::set_comparison::{ComparisonOp, ReductionOp};
 use substrait::proto::expression::subquery::set_predicate::PredicateOp;
 use substrait::proto::expression::subquery::SubqueryType;
 
@@ -94,8 +95,53 @@ pub async fn from_subquery(
                     ),
                 }
             }
-            other_type => {
-                substrait_err!("Subquery type {other_type:?} not implemented")
+            SubqueryType::SetComparison(comparison) => {
+                let left = comparison.left.as_ref().ok_or_else(|| {
+                    substrait_datafusion_err!("SetComparison requires a left expression")
+                })?;
+                let right = comparison.right.as_ref().ok_or_else(|| {
+                    substrait_datafusion_err!("SetComparison requires a right relation")
+                })?;
+                let reduction_op = match ReductionOp::try_from(comparison.reduction_op) {
+                    Ok(ReductionOp::Any) => SetQuantifier::Any,
+                    Ok(ReductionOp::All) => SetQuantifier::All,
+                    _ => {
+                        return substrait_err!(
+                            "Unsupported reduction op for SetComparison: {}",
+                            comparison.reduction_op
+                        )
+                    }
+                };
+                let comparison_op = match ComparisonOp::try_from(comparison.comparison_op)
+                {
+                    Ok(ComparisonOp::Eq) => Operator::Eq,
+                    Ok(ComparisonOp::Ne) => Operator::NotEq,
+                    Ok(ComparisonOp::Lt) => Operator::Lt,
+                    Ok(ComparisonOp::Gt) => Operator::Gt,
+                    Ok(ComparisonOp::Le) => Operator::LtEq,
+                    Ok(ComparisonOp::Ge) => Operator::GtEq,
+                    _ => {
+                        return substrait_err!(
+                            "Unsupported comparison op for SetComparison: {}",
+                            comparison.comparison_op
+                        )
+                    }
+                };
+
+                let left_expr = consumer.consume_expression(left, input_schema).await?;
+                let plan = consumer.consume_rel(right).await?;
+                let outer_ref_columns = plan.all_out_ref_exprs();
+
+                Ok(Expr::SetComparison(SetComparison::new(
+                    Box::new(left_expr),
+                    Subquery {
+                        subquery: Arc::new(plan),
+                        outer_ref_columns,
+                        spans: Spans::new(),
+                    },
+                    comparison_op,
+                    reduction_op,
+                )))
             }
         },
         None => {

datafusion/substrait/src/logical_plan/producer/expr/mod.rs

@@ -141,9 +141,7 @@ pub fn to_substrait_rex(
         Expr::InList(expr) => producer.handle_in_list(expr, schema),
         Expr::Exists(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"),
         Expr::InSubquery(expr) => producer.handle_in_subquery(expr, schema),
-        Expr::SetComparison(expr) => {
-            not_impl_err!("Cannot convert {expr:?} to Substrait")
-        }
+        Expr::SetComparison(expr) => producer.handle_set_comparison(expr, schema),
         Expr::ScalarSubquery(expr) => {
             not_impl_err!("Cannot convert {expr:?} to Substrait")
         }

datafusion/substrait/src/logical_plan/producer/expr/subquery.rs

@@ -16,8 +16,10 @@
 // under the License.
 
 use crate::logical_plan::producer::SubstraitProducer;
-use datafusion::common::DFSchemaRef;
-use datafusion::logical_expr::expr::InSubquery;
+use datafusion::common::{substrait_err, DFSchemaRef};
+use datafusion::logical_expr::expr::{InSubquery, SetComparison, SetQuantifier};
+use datafusion::logical_expr::Operator;
+use substrait::proto::expression::subquery::set_comparison::{ComparisonOp, ReductionOp};
 use substrait::proto::expression::subquery::InPredicate;
 use substrait::proto::expression::{RexType, ScalarFunction};
 use substrait::proto::function_argument::ArgType;
@@ -70,3 +72,53 @@ pub fn from_in_subquery(
         Ok(substrait_subquery)
     }
 }
+
+fn comparison_op_to_proto(op: &Operator) -> datafusion::common::Result<ComparisonOp> {
+    match op {
+        Operator::Eq => Ok(ComparisonOp::Eq),
+        Operator::NotEq => Ok(ComparisonOp::Ne),
+        Operator::Lt => Ok(ComparisonOp::Lt),
+        Operator::Gt => Ok(ComparisonOp::Gt),
+        Operator::LtEq => Ok(ComparisonOp::Le),
+        Operator::GtEq => Ok(ComparisonOp::Ge),
+        _ => substrait_err!("Unsupported operator {op:?} for SetComparison subquery"),
+    }
+}
+
+fn reduction_op_to_proto(
+    quantifier: &SetQuantifier,
+) -> datafusion::common::Result<ReductionOp> {
+    match quantifier {
+        SetQuantifier::Any => Ok(ReductionOp::Any),
+        SetQuantifier::All => Ok(ReductionOp::All),
+    }
+}
+
+pub fn from_set_comparison(
+    producer: &mut impl SubstraitProducer,
+    set_comparison: &SetComparison,
+    schema: &DFSchemaRef,
+) -> datafusion::common::Result<Expression> {
+    let comparison_op = comparison_op_to_proto(&set_comparison.op)? as i32;
+    let reduction_op = reduction_op_to_proto(&set_comparison.quantifier)? as i32;
+    let left = producer.handle_expr(set_comparison.expr.as_ref(), schema)?;
+    let subquery_plan =
+        producer.handle_plan(set_comparison.subquery.subquery.as_ref())?;
+
+    Ok(Expression {
+        rex_type: Some(RexType::Subquery(Box::new(
+            substrait::proto::expression::Subquery {
+                subquery_type: Some(
+                    substrait::proto::expression::subquery::SubqueryType::SetComparison(
+                        Box::new(substrait::proto::expression::subquery::SetComparison {
+                            reduction_op,
+                            comparison_op,
+                            left: Some(Box::new(left)),
+                            right: Some(subquery_plan),
+                        }),
+                    ),
+                ),
+            },
+        ))),
+    })
+}

datafusion/substrait/src/logical_plan/producer/substrait_producer.rs

@@ -20,14 +20,17 @@ use crate::logical_plan::producer::{
     from_aggregate, from_aggregate_function, from_alias, from_between, from_binary_expr,
     from_case, from_cast, from_column, from_distinct, from_empty_relation, from_filter,
     from_in_list, from_in_subquery, from_join, from_like, from_limit, from_literal,
-    from_projection, from_repartition, from_scalar_function, from_sort,
-    from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr, from_union,
-    from_values, from_window, from_window_function, to_substrait_rel, to_substrait_rex,
+    from_projection, from_repartition, from_scalar_function, from_set_comparison,
+    from_sort, from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr,
+    from_union, from_values, from_window, from_window_function, to_substrait_rel,
+    to_substrait_rex,
 };
 use datafusion::common::{substrait_err, Column, DFSchemaRef, ScalarValue};
 use datafusion::execution::registry::SerializerRegistry;
 use datafusion::execution::SessionState;
-use datafusion::logical_expr::expr::{Alias, InList, InSubquery, WindowFunction};
+use datafusion::logical_expr::expr::{
+    Alias, InList, InSubquery, SetComparison, WindowFunction,
+};
 use datafusion::logical_expr::{
     expr, Aggregate, Between, BinaryExpr, Case, Cast, Distinct, EmptyRelation, Expr,
     Extension, Filter, Join, Like, Limit, LogicalPlan, Projection, Repartition, Sort,
@@ -359,6 +362,14 @@ pub trait SubstraitProducer: Send + Sync + Sized {
     ) -> datafusion::common::Result<Expression> {
         from_in_subquery(self, in_subquery, schema)
     }
+
+    fn handle_set_comparison(
+        &mut self,
+        set_comparison: &SetComparison,
+        schema: &DFSchemaRef,
+    ) -> datafusion::common::Result<Expression> {
+        from_set_comparison(self, set_comparison, schema)
+    }
 }
 
 pub struct DefaultSubstraitProducer<'a> {

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

@@ -29,14 +29,15 @@ use std::mem::size_of_val;
 
 use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit};
 use datafusion::common::tree_node::Transformed;
-use datafusion::common::{not_impl_err, plan_err, DFSchema, DFSchemaRef};
+use datafusion::common::{not_impl_err, plan_err, DFSchema, DFSchemaRef, Spans};
 use datafusion::error::Result;
 use datafusion::execution::registry::SerializerRegistry;
 use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::execution::session_state::SessionStateBuilder;
+use datafusion::logical_expr::expr::{SetComparison, SetQuantifier};
 use datafusion::logical_expr::{
-    EmptyRelation, Extension, InvariantLevel, LogicalPlan, PartitionEvaluator,
-    Repartition, UserDefinedLogicalNode, Values, Volatility,
+    EmptyRelation, Extension, InvariantLevel, LogicalPlan, Operator, PartitionEvaluator,
+    Repartition, Subquery, UserDefinedLogicalNode, Values, Volatility,
 };
 use datafusion::optimizer::simplify_expressions::expr_simplifier::THRESHOLD_INLINE_INLIST;
 use datafusion::prelude::*;
@@ -689,6 +690,29 @@ async fn roundtrip_exists_filter() -> Result<()> {
     Ok(())
 }
 
+// assemble logical plan manually to ensure SetComparison expr is present (not rewrite away)
+#[tokio::test]
+async fn roundtrip_set_comparison_any_substrait() -> Result<()> {
+    let ctx = create_context().await?;
+    let plan = build_set_comparison_plan(&ctx, SetQuantifier::Any, Operator::Gt).await?;
+    let proto = to_substrait_plan(&plan, &ctx.state())?;
+    let roundtrip_plan = from_substrait_plan(&ctx.state(), &proto).await?;
+    assert_set_comparison_predicate(&roundtrip_plan, Operator::Gt, SetQuantifier::Any);
+    Ok(())
+}
+
+// assemble logical plan manually to ensure SetComparison expr is present (not rewrite away)
+#[tokio::test]
+async fn roundtrip_set_comparison_all_substrait() -> Result<()> {
+    let ctx = create_context().await?;
+    let plan =
+        build_set_comparison_plan(&ctx, SetQuantifier::All, Operator::NotEq).await?;
+    let proto = to_substrait_plan(&plan, &ctx.state())?;
+    let roundtrip_plan = from_substrait_plan(&ctx.state(), &proto).await?;
+    assert_set_comparison_predicate(&roundtrip_plan, Operator::NotEq, SetQuantifier::All);
+    Ok(())
+}
+
 #[tokio::test]
 async fn roundtrip_not_exists_filter_left_anti_join() -> Result<()> {
     let plan = generate_plan_from_sql(
@@ -1865,6 +1889,56 @@ async fn assert_substrait_sql(substrait_plan: Plan, sql: &str) -> Result<()> {
     Ok(())
 }
 
+async fn build_set_comparison_plan(
+    ctx: &SessionContext,
+    quantifier: SetQuantifier,
+    op: Operator,
+) -> Result<LogicalPlan> {
+    let base_scan = ctx.table("data").await?.into_unoptimized_plan();
+    let subquery_scan = ctx.table("data2").await?.into_unoptimized_plan();
+    let subquery_plan = LogicalPlanBuilder::from(subquery_scan)
+        .project(vec![col("data2.a")])?
+        .build()?;
+    let predicate = Expr::SetComparison(SetComparison::new(
+        Box::new(col("data.a")),
+        Subquery {
+            subquery: Arc::new(subquery_plan),
+            outer_ref_columns: vec![],
+            spans: Spans::new(),
+        },
+        op,
+        quantifier,
+    ));
+
+    LogicalPlanBuilder::from(base_scan)
+        .filter(predicate)?
+        .project(vec![col("data.a")])?
+        .build()
+}
+
+fn assert_set_comparison_predicate(
+    plan: &LogicalPlan,
+    expected_op: Operator,
+    expected_quantifier: SetQuantifier,
+) {
+    let predicate = match plan {
+        LogicalPlan::Projection(p) => match p.input.as_ref() {
+            LogicalPlan::Filter(filter) => &filter.predicate,
+            other => panic!("expected Filter inside Projection, got {other:?}"),
+        },
+        LogicalPlan::Filter(filter) => &filter.predicate,
+        other => panic!("expected Filter plan, got {other:?}"),
+    };
+
+    match predicate {
+        Expr::SetComparison(set_comparison) => {
+            assert_eq!(set_comparison.op, expected_op);
+            assert_eq!(set_comparison.quantifier, expected_quantifier);
+        }
+        other => panic!("expected SetComparison predicate, got {other:?}"),
+    }
+}
+
 async fn roundtrip_fill_na(sql: &str) -> Result<()> {
     let ctx = create_context().await?;
     let df = ctx.sql(sql).await?;