diff --git a/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java b/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java index 0f3b37fad..cc4294efb 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java @@ -246,31 +246,31 @@ public SqlKind visit(NamedUpdate update, EmptyVisitationContext context) throws @Override public SqlKind visit(ScatterExchange exchange, EmptyVisitationContext context) throws RuntimeException { - return SqlKind.OTHER_DDL; + return QUERY_KIND; } @Override public SqlKind visit(SingleBucketExchange exchange, EmptyVisitationContext context) throws RuntimeException { - return SqlKind.OTHER_DDL; + return QUERY_KIND; } @Override public SqlKind visit(MultiBucketExchange exchange, EmptyVisitationContext context) throws RuntimeException { - return SqlKind.OTHER_DDL; + return QUERY_KIND; } @Override public SqlKind visit(RoundRobinExchange exchange, EmptyVisitationContext context) throws RuntimeException { - return SqlKind.OTHER_DDL; + return QUERY_KIND; } @Override public SqlKind visit(BroadcastExchange exchange, EmptyVisitationContext context) throws RuntimeException { - return SqlKind.OTHER_DDL; + return QUERY_KIND; } @Override diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java index f6a221f8b..c05f5b6f3 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java @@ -4,6 +4,7 @@ import com.google.common.collect.ImmutableSet; import io.substrait.expression.Expression; import io.substrait.expression.Expression.SortDirection; +import io.substrait.expression.FieldReference; import io.substrait.expression.FunctionArg; import io.substrait.expression.MaskExpression; import io.substrait.extension.FunctionBindingResolver; @@ -73,6 +74,7 @@ import org.apache.calcite.prepare.Prepare; import org.apache.calcite.rel.RelCollation; import org.apache.calcite.rel.RelCollations; +import org.apache.calcite.rel.RelDistributions; import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.AggregateCall; @@ -994,7 +996,28 @@ public RelNode visit(NamedUpdate update, Context context) { @Override public RelNode visit(ScatterExchange exchange, Context context) throws RuntimeException { - return visitFallback(exchange, context); + RelNode child = exchange.getInput().accept(this, context); + List keys = + exchange.getFields().stream() + .map(SubstraitRelNodeConverter::topLevelFieldIndex) + .collect(Collectors.toList()); + // Calcite sorts hash keys, but Substrait hashes in field order, which other exchanges may rely + // on to line up partitions. + if (!keys.equals(keys.stream().sorted().collect(Collectors.toList()))) { + throw new UnsupportedOperationException( + "A scatter exchange whose fields are not in ascending order is not supported"); + } + RelNode node = relBuilder.push(child).exchange(RelDistributions.hash(keys)).build(); + return applyRelCommon(node, exchange, child); + } + + private static int topLevelFieldIndex(FieldReference field) { + if (!field.isSimpleRootReference() + || !(field.segments().get(0) instanceof FieldReference.StructField)) { + throw new UnsupportedOperationException( + "A scatter exchange on anything other than a top-level field is not supported"); + } + return ((FieldReference.StructField) field.segments().get(0)).offset(); } @Override @@ -1009,12 +1032,17 @@ public RelNode visit(MultiBucketExchange exchange, Context context) throws Runti @Override public RelNode visit(RoundRobinExchange exchange, Context context) throws RuntimeException { - return visitFallback(exchange, context); + RelNode child = exchange.getInput().accept(this, context); + RelNode node = + relBuilder.push(child).exchange(RelDistributions.ROUND_ROBIN_DISTRIBUTED).build(); + return applyRelCommon(node, exchange, child); } @Override public RelNode visit(BroadcastExchange exchange, Context context) throws RuntimeException { - return visitFallback(exchange, context); + RelNode child = exchange.getInput().accept(this, context); + RelNode node = relBuilder.push(child).exchange(RelDistributions.BROADCAST_DISTRIBUTED).build(); + return applyRelCommon(node, exchange, child); } @Override diff --git a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java index 257aeaacd..2abe55c67 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java @@ -36,6 +36,10 @@ import io.substrait.relation.Set; import io.substrait.relation.Sort; import io.substrait.relation.VirtualTableScan; +import io.substrait.relation.physical.BroadcastExchange; +import io.substrait.relation.physical.RoundRobinExchange; +import io.substrait.relation.physical.ScatterExchange; +import io.substrait.relation.physical.SingleBucketExchange; import io.substrait.type.NamedStruct; import io.substrait.type.Type; import io.substrait.type.TypeCreator; @@ -49,6 +53,7 @@ import java.util.stream.Stream; import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelDistribution; import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelFieldCollation.Direction; import org.apache.calcite.rel.RelNode; @@ -797,7 +802,40 @@ private static Expression.SortDirection asSortDirection(RelFieldCollation collat */ @Override public Rel visit(org.apache.calcite.rel.core.Exchange exchange) { - return super.visit(exchange); + // A SortExchange also sorts each partition, which no Substrait exchange expresses. + if (exchange instanceof org.apache.calcite.rel.core.SortExchange) { + throw new UnsupportedOperationException("SortExchange is not supported"); + } + Rel input = apply(exchange.getInput()); + RelDistribution distribution = exchange.getDistribution(); + // A Calcite plan carries no partition count; 0 leaves it to the consumer. + switch (distribution.getType()) { + case HASH_DISTRIBUTED: + return ScatterExchange.builder() + .input(input) + .addAllFields( + distribution.getKeys().stream() + .map(key -> FieldReference.newInputRelReference(key, input)) + .collect(Collectors.toList())) + .partitionCount(0) + .build(); + case BROADCAST_DISTRIBUTED: + return BroadcastExchange.builder().input(input).partitionCount(0).build(); + case ROUND_ROBIN_DISTRIBUTED: + // Calcite assigns records to instances in turn, i.e. per record. + return RoundRobinExchange.builder().input(input).exact(true).partitionCount(0).build(); + case RANDOM_DISTRIBUTED: + return RoundRobinExchange.builder().input(input).exact(false).partitionCount(0).build(); + case SINGLETON: + return SingleBucketExchange.builder() + .input(input) + .expression(ExpressionCreator.i32(false, 0)) + .partitionCount(1) + .build(); + default: + throw new UnsupportedOperationException( + "Unsupported exchange distribution: " + distribution.getType()); + } } /** diff --git a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java index 70c5807dd..73cca2adf 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java @@ -13,6 +13,7 @@ import io.substrait.expression.AggregateFunctionInvocation; import io.substrait.expression.EnumArg; import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; import io.substrait.expression.FunctionArg; import io.substrait.expression.FunctionOption; import io.substrait.expression.ImmutableAggregateFunctionInvocation; @@ -25,6 +26,12 @@ import io.substrait.relation.Join.JoinType; import io.substrait.relation.Rel; import io.substrait.relation.Set.SetOp; +import io.substrait.relation.physical.AbstractExchangeRel; +import io.substrait.relation.physical.BroadcastExchange; +import io.substrait.relation.physical.RoundRobinExchange; +import io.substrait.relation.physical.ScatterExchange; +import io.substrait.relation.physical.SingleBucketExchange; +import io.substrait.relation.physical.TargetType; import io.substrait.type.Type; import java.util.ArrayList; import java.util.List; @@ -35,6 +42,9 @@ import org.apache.calcite.plan.hep.HepPlanner; import org.apache.calcite.plan.hep.HepProgram; import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.rel.RelCollations; +import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelDistributions; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.AggregateCall; import org.apache.calcite.rel.rules.CoreRules; @@ -1425,4 +1435,166 @@ void rejectsCustomComparisonFunction() { assertThrows(UnsupportedOperationException.class, () -> substraitToCalcite.convert(sort)); } } + + @Nested + class Exchange { + @Test + void scatterRoundTrip() { + assertFullRoundTrip( + ScatterExchange.builder() + .input(commonTable) + .addFields(sb.fieldReference(commonTable, 0), sb.fieldReference(commonTable, 2)) + .partitionCount(0) + .build()); + } + + @Test + void broadcastRoundTrip() { + assertFullRoundTrip(BroadcastExchange.builder().input(commonTable).partitionCount(0).build()); + } + + @Test + void exactRoundRobinRoundTrip() { + assertFullRoundTrip( + RoundRobinExchange.builder().input(commonTable).exact(true).partitionCount(0).build()); + } + + @Test + void approximateRoundRobinConvertsToRoundRobin() { + Rel exchange = + RoundRobinExchange.builder().input(commonTable).exact(false).partitionCount(0).build(); + + assertEquals( + RelDistribution.Type.ROUND_ROBIN_DISTRIBUTED, + distributionType(substraitToCalcite.convert(exchange))); + } + + @Test + void partitionCountAndTargetsAreDropped() { + Rel exchange = + BroadcastExchange.builder() + .input(commonTable) + .partitionCount(4) + .addTargets( + AbstractExchangeRel.ExchangeTarget.builder() + .addPartitionIds(0) + .type(TargetType.Uri.builder().uri("worker-0").build()) + .build()) + .build(); + + assertEquals( + RelDistribution.Type.BROADCAST_DISTRIBUTED, + distributionType(substraitToCalcite.convert(exchange))); + } + + @Test + void rejectsUnorderedScatterFields() { + Rel exchange = + ScatterExchange.builder() + .input(commonTable) + .addFields(sb.fieldReference(commonTable, 2), sb.fieldReference(commonTable, 0)) + .partitionCount(0) + .build(); + + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, () -> substraitToCalcite.convert(exchange)); + assertEquals( + "A scatter exchange whose fields are not in ascending order is not supported", + e.getMessage()); + } + + @Test + void rejectsNestedScatterField() { + Rel input = sb.namedScan(List.of("example"), List.of("s", "x"), List.of(R.struct(R.I32))); + Rel exchange = + ScatterExchange.builder() + .input(input) + .addFields(sb.fieldReference(input, 0).dereferenceStruct(0)) + .partitionCount(0) + .build(); + + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, () -> substraitToCalcite.convert(exchange)); + assertEquals( + "A scatter exchange on anything other than a top-level field is not supported", + e.getMessage()); + } + + @Test + void randomDistributionConvertsToApproximateRoundRobin() { + RelNode exchange = + builder + .values(new String[] {"a"}, 1) + .exchange(RelDistributions.RANDOM_DISTRIBUTED) + .build(); + + RoundRobinExchange converted = + assertInstanceOf( + RoundRobinExchange.class, SubstraitRelVisitor.convert(exchange, converterProvider)); + assertFalse(converted.getExact()); + assertEquals(0, converted.getPartitionCount()); + } + + @Test + void singletonConvertsToSingleBucket() { + RelNode exchange = + builder.values(new String[] {"a"}, 1).exchange(RelDistributions.SINGLETON).build(); + + SingleBucketExchange converted = + assertInstanceOf( + SingleBucketExchange.class, SubstraitRelVisitor.convert(exchange, converterProvider)); + assertEquals(ExpressionCreator.i32(false, 0), converted.getExpression()); + assertEquals(1, converted.getPartitionCount()); + } + + @Test + void rejectsRangeDistribution() { + RelNode exchange = + builder + .values(new String[] {"a"}, 1) + .exchange(RelDistributions.range(List.of(0))) + .build(); + + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, + () -> SubstraitRelVisitor.convert(exchange, converterProvider)); + assertEquals("Unsupported exchange distribution: RANGE_DISTRIBUTED", e.getMessage()); + } + + @Test + void rejectsSortExchange() { + RelNode exchange = + builder + .values(new String[] {"a"}, 1) + .sortExchange(RelDistributions.hash(List.of(0)), RelCollations.of(0)) + .build(); + + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, + () -> SubstraitRelVisitor.convert(exchange, converterProvider)); + assertEquals("SortExchange is not supported", e.getMessage()); + } + + @Test + void rootKeepsOutputNames() { + Plan.Root root = + sb.root( + BroadcastExchange.builder().input(commonTable).partitionCount(0).build(), + List.of("w", "x", "y", "z")); + + RelNode projected = substraitToCalcite.convert(root).project(true); + + assertEquals(List.of("w", "x", "y", "z"), projected.getRowType().getFieldNames()); + } + + private RelDistribution.Type distributionType(RelNode node) { + return assertInstanceOf(org.apache.calcite.rel.core.Exchange.class, node) + .getDistribution() + .getType(); + } + } }