Skip to content

feat(isthmus): convert exchange relations to and from Calcite #1426

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anasik wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(isthmus): convert exchange relations to and from Calcite #1426

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions isthmus/src/main/java/io/substrait/isthmus/SqlKindFromRel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
34 changes: 31 additions & 3 deletions isthmus/src/main/java/io/substrait/isthmus/SubstraitRelNodeConverter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<Integer> 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
Expand All @@ -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
Expand Down
40 changes: 39 additions & 1 deletion isthmus/src/main/java/io/substrait/isthmus/SubstraitRelVisitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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());
}
}

/**
Expand Down
172 changes: 172 additions & 0 deletions isthmus/src/test/java/io/substrait/isthmus/SubstraitRelNodeConverterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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();
}
}
}
Loading