Skip to content

Commit 91d2886

Browse files
authored
fix(substrait): disallow union with a single input (#13023)
* fix(substrait): disallow union with a single input * flip if condition
1 parent 227908f commit 91d2886

File tree

1 file changed

+34
-64
lines changed

1 file changed

+34
-64
lines changed

datafusion/substrait/src/logical_plan/consumer.rs

Lines changed: 34 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -937,72 +937,42 @@ pub async fn from_substrait_rel(
937937
}
938938
}
939939
Some(RelType::Set(set)) => match set_rel::SetOp::try_from(set.op) {
940-
Ok(set_op) => match set_op {
941-
set_rel::SetOp::UnionAll => {
942-
if !set.inputs.is_empty() {
943-
union_rels(&set.inputs, ctx, extensions, true).await
944-
} else {
945-
not_impl_err!("Union relation requires at least one input")
946-
}
947-
}
948-
set_rel::SetOp::UnionDistinct => {
949-
if !set.inputs.is_empty() {
950-
union_rels(&set.inputs, ctx, extensions, false).await
951-
} else {
952-
not_impl_err!("Union relation requires at least one input")
953-
}
954-
}
955-
set_rel::SetOp::IntersectionPrimary => {
956-
if set.inputs.len() >= 2 {
957-
LogicalPlanBuilder::intersect(
958-
from_substrait_rel(ctx, &set.inputs[0], extensions).await?,
959-
union_rels(&set.inputs[1..], ctx, extensions, true).await?,
960-
false,
961-
)
962-
} else {
963-
not_impl_err!(
964-
"Primary Intersect relation requires at least two inputs"
965-
)
966-
}
967-
}
968-
set_rel::SetOp::IntersectionMultiset => {
969-
if set.inputs.len() >= 2 {
970-
intersect_rels(&set.inputs, ctx, extensions, false).await
971-
} else {
972-
not_impl_err!(
973-
"Multiset Intersect relation requires at least two inputs"
974-
)
975-
}
976-
}
977-
set_rel::SetOp::IntersectionMultisetAll => {
978-
if set.inputs.len() >= 2 {
979-
intersect_rels(&set.inputs, ctx, extensions, true).await
980-
} else {
981-
not_impl_err!(
982-
"MultisetAll Intersect relation requires at least two inputs"
983-
)
984-
}
985-
}
986-
set_rel::SetOp::MinusPrimary => {
987-
if set.inputs.len() >= 2 {
988-
except_rels(&set.inputs, ctx, extensions, false).await
989-
} else {
990-
not_impl_err!(
991-
"Primary Minus relation requires at least two inputs"
992-
)
993-
}
994-
}
995-
set_rel::SetOp::MinusPrimaryAll => {
996-
if set.inputs.len() >= 2 {
997-
except_rels(&set.inputs, ctx, extensions, true).await
998-
} else {
999-
not_impl_err!(
1000-
"PrimaryAll Minus relation requires at least two inputs"
1001-
)
940+
Ok(set_op) => {
941+
if set.inputs.len() < 2 {
942+
substrait_err!("Set operation requires at least two inputs")
943+
} else {
944+
match set_op {
945+
set_rel::SetOp::UnionAll => {
946+
union_rels(&set.inputs, ctx, extensions, true).await
947+
}
948+
set_rel::SetOp::UnionDistinct => {
949+
union_rels(&set.inputs, ctx, extensions, false).await
950+
}
951+
set_rel::SetOp::IntersectionPrimary => {
952+
LogicalPlanBuilder::intersect(
953+
from_substrait_rel(ctx, &set.inputs[0], extensions)
954+
.await?,
955+
union_rels(&set.inputs[1..], ctx, extensions, true)
956+
.await?,
957+
false,
958+
)
959+
}
960+
set_rel::SetOp::IntersectionMultiset => {
961+
intersect_rels(&set.inputs, ctx, extensions, false).await
962+
}
963+
set_rel::SetOp::IntersectionMultisetAll => {
964+
intersect_rels(&set.inputs, ctx, extensions, true).await
965+
}
966+
set_rel::SetOp::MinusPrimary => {
967+
except_rels(&set.inputs, ctx, extensions, false).await
968+
}
969+
set_rel::SetOp::MinusPrimaryAll => {
970+
except_rels(&set.inputs, ctx, extensions, true).await
971+
}
972+
_ => not_impl_err!("Unsupported set operator: {set_op:?}"),
1002973
}
1003974
}
1004-
_ => not_impl_err!("Unsupported set operator: {set_op:?}"),
1005-
},
975+
}
1006976
Err(e) => not_impl_err!("Invalid set operation type {}: {e}", set.op),
1007977
},
1008978
Some(RelType::ExtensionLeaf(extension)) => {

0 commit comments

Comments
 (0)