diff --git a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestSchedulingUtils.java b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestSchedulingUtils.java index baeeb450f88d..9f93ae3a4091 100644 --- a/core/trino-main/src/test/java/io/trino/execution/scheduler/TestSchedulingUtils.java +++ b/core/trino-main/src/test/java/io/trino/execution/scheduler/TestSchedulingUtils.java @@ -72,6 +72,34 @@ public void testCanStreamNoJoin() assertThat(SchedulingUtils.canStream(parentSubPlan, valuesSubPlan("a"))).isTrue(); } + @Test + public void testCanStreamUnion() + { + /* + parent(union) + / \ \ + -------------------------------------- stage boundary + a b c + */ + SubPlan aSubPlan = valuesSubPlan("a"); + RemoteSourceNode remoteSourceA = remoteSource("a"); + + SubPlan bSubPlan = valuesSubPlan("b"); + RemoteSourceNode remoteSourceB = remoteSource("b"); + + SubPlan cSubPlan = valuesSubPlan("c"); + RemoteSourceNode remoteSourceC = remoteSource("c"); + + SubPlan parentSubPlan = createSubPlan( + "parent", + union("union", ImmutableList.of(remoteSourceA, remoteSourceB, remoteSourceC)), + ImmutableList.of(aSubPlan, bSubPlan, cSubPlan)); + + assertThat(SchedulingUtils.canStream(parentSubPlan, aSubPlan)).isTrue(); + assertThat(SchedulingUtils.canStream(parentSubPlan, bSubPlan)).isTrue(); + assertThat(SchedulingUtils.canStream(parentSubPlan, cSubPlan)).isTrue(); + } + @Test public void testCanStreamJoin() {