Skip to content

feat(isthmus): observe field reference types #1080

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
alexandrefimov wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(isthmus): observe field reference types #1080

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
95 changes: 92 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 @@ -320,6 +320,7 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
}

RelNode child = aggregate.getInput().accept(this, context);
context.enterScope(AnchoredInput.of(aggregate.getInput().getRelAnchor(), child.getRowType()));
List<List<RexNode>> groupExprLists =
aggregate.getGroupings().stream()
.map(
Expand Down Expand Up @@ -383,6 +384,8 @@ public RelNode visit(Aggregate aggregate, Context context) throws RuntimeExcepti
}
}

exitUncorrelatedScope(context, "aggregate");

// RelBuilder deduplicates equal aggregate calls, and AggregateCall equality ignores the stored
// type: two measures of the same function that differ only by their declared output type would
// collapse into one column. Opt out of deduplication for exactly those aggregates — narrowly,
Expand Down Expand Up @@ -605,10 +608,12 @@ private AggregateCall fromMeasure(
@Override
public RelNode visit(Sort sort, Context context) throws RuntimeException {
RelNode child = sort.getInput().accept(this, context);
context.enterScope(AnchoredInput.of(sort.getInput().getRelAnchor(), child.getRowType()));
List<RexNode> sortExpressions =
sort.getSortFields().stream()
.map(sortField -> directedRexNode(sortField, context))
.collect(Collectors.toList());
exitUncorrelatedScope(context, "sort");
RelNode node = relBuilder.push(child).sort(sortExpressions).build();
return applyRemap(node, sort.getRemap());
}
Expand Down Expand Up @@ -641,13 +646,15 @@ private RexNode directedRexNode(Expression.SortField sortField, Context context)
@Override
public RelNode visit(Fetch fetch, Context context) throws RuntimeException {
RelNode child = fetch.getInput().accept(this, context);
context.enterScope(AnchoredInput.of(fetch.getInput().getRelAnchor(), child.getRowType()));
// Offset/count are expressions; pass them through to Calcite as RexNodes so non-literal (e.g.
// dynamic-parameter) offset/count are preserved. An unset offset means 0 and an unset count
// means LIMIT ALL.
RexNode offset =
fetch.getOffset().map(e -> e.accept(expressionRexConverter, context)).orElse(null);
RexNode count =
fetch.getCount().map(e -> e.accept(expressionRexConverter, context)).orElse(null);
exitUncorrelatedScope(context, "fetch");
RelNode node = relBuilder.push(child).sortLimit(offset, count, ImmutableList.of()).build();
return applyRemap(node, fetch.getRemap());
}
Expand Down Expand Up @@ -688,9 +695,8 @@ private RelFieldCollation toRelFieldCollation(Expression.SortField sortField, Co
@Override
public RelNode visit(NamedUpdate update, Context context) {
relBuilder.scan(update.getNames());
context.enterScope(AnchoredInput.of(Optional.empty(), relBuilder.peek().getRowType()));
RexNode condition = update.getCondition().accept(expressionRexConverter, context);
relBuilder.filter(condition);
RelNode inputForModify = relBuilder.build();

NamedStruct tableSchema = update.getTableSchema();
List<String> fieldNames = tableSchema.names();
Expand All @@ -705,6 +711,10 @@ public RelNode visit(NamedUpdate update, Context context) {
transform.getTransformation().accept(expressionRexConverter, context));
}

java.util.Set<CorrelationId> correlationIds = context.exitScope();
relBuilder.filter(correlationIds, condition);
RelNode inputForModify = relBuilder.build();

assert relBuilder.getRelOptSchema() != null;
final RelOptTable table = relBuilder.getRelOptSchema().getTableForMember(update.getNames());

Expand Down Expand Up @@ -929,6 +939,23 @@ protected RelNode applyRemap(RelNode relNode, Optional<Rel.Remap> remap) {
return relNode;
}

/**
* Exits a scope entered only to resolve input field types, on an operator that cannot carry
* correlation variables. Calcite's Aggregate, Sort and Fetch have no {@code variablesSet}, so a
* correlation resolved against such an operator's input has nowhere to be declared and is
* rejected rather than silently dropped.
*
* @param context the conversion context whose innermost scope is exited
* @param relName the relation kind, used in the failure message
*/
private static void exitUncorrelatedScope(Context context, String relName) {
java.util.Set<CorrelationId> correlationIds = context.exitScope();
if (!correlationIds.isEmpty()) {
throw new UnsupportedOperationException(
"Outer references bound to the " + relName + " input are not supported");
}
}

private RelNode applyRemap(RelNode relNode, Rel.Remap remap) {
RelDataType rowType = relNode.getRowType();
List<String> fieldNames = rowType.getFieldNames();
Expand Down Expand Up @@ -987,6 +1014,9 @@ public static class Context implements VisitationContext {
/** Maps a {@code rel_anchor} to the single {@link CorrelationId} minted for it. */
private final Map<Integer, CorrelationId> correlationIdByAnchor = new HashMap<>();

/** Lambda parameter types by nesting level, innermost on top. */
private final Deque<List<RelDataType>> lambdaParameterTypes = new ArrayDeque<>();

/**
* Every {@code rel_anchor} that has entered a scope. Resolution keys {@link #scopeByAnchor} and
* {@link #correlationIdByAnchor} purely by anchor value, which is only sound if anchors are
Expand All @@ -998,6 +1028,7 @@ public static class Context implements VisitationContext {

/** One correlation scope per enclosing relational operator. */
private static final class Scope {
final List<RelDataType> inputRowTypes = new ArrayList<>();
final Map<Integer, RelDataType> rowTypeByAnchor = new HashMap<>();
final java.util.Set<CorrelationId> correlationIds = new HashSet<>();
}
Expand All @@ -1013,13 +1044,14 @@ public static Context newContext() {

/**
* Enters a correlation scope for a relational operator, recording the {@code rel_anchor} (if
* any) carried by each of its inputs.
* any) and row type carried by each of its inputs.
*
* @param inputs the operator's inputs paired with their anchors
*/
public void enterScope(final AnchoredInput... inputs) {
final Scope scope = new Scope();
for (final AnchoredInput input : inputs) {
scope.inputRowTypes.add(input.rowType);
if (input.anchor.isPresent()) {
final int anchor = input.anchor.get();
if (!seenAnchors.add(anchor)) {
Expand All @@ -1036,6 +1068,63 @@ public void enterScope(final AnchoredInput... inputs) {
scopes.push(scope);
}

/**
* Returns the type of a field in the current operator's flattened input row.
*
* @param fieldIndex zero-based field index across all current inputs
* @return the Calcite field type from the input relation
* @throws IllegalStateException if expression conversion has no current relational input
* @throws IndexOutOfBoundsException if the field index is outside the flattened input row
*/
public RelDataType getInputFieldType(final int fieldIndex) {
if (scopes.isEmpty()) {
throw new IllegalStateException("No input row type is available for field reference");
}

int remainingIndex = fieldIndex;
int fieldCount = 0;
for (final RelDataType inputRowType : scopes.peek().inputRowTypes) {
final int inputFieldCount = inputRowType.getFieldCount();
fieldCount += inputFieldCount;
if (remainingIndex < inputFieldCount) {
return inputRowType.getFieldList().get(remainingIndex).getType();
}
remainingIndex -= inputFieldCount;
}
throw new IndexOutOfBoundsException(
"Field index " + fieldIndex + " is outside input row with " + fieldCount + " fields");
}

/**
* Enters a lambda scope with its Calcite parameter types.
*
* @param parameterTypes parameter types in declaration order
*/
public void enterLambdaScope(final List<RelDataType> parameterTypes) {
lambdaParameterTypes.push(List.copyOf(parameterTypes));
}

/** Exits the innermost lambda scope. */
public void exitLambdaScope() {
lambdaParameterTypes.pop();
}

/**
* Returns a parameter type from the innermost lambda scope.
*
* @param parameterIndex zero-based parameter index
* @return the Calcite parameter type
* @throws IllegalStateException if expression conversion has no current lambda
* @throws IndexOutOfBoundsException if the index is outside the lambda's parameter list
*/
public RelDataType getLambdaParameterType(final int parameterIndex) {
if (lambdaParameterTypes.isEmpty()) {
throw new IllegalStateException(
"No lambda parameter type is available for field reference");
}
return lambdaParameterTypes.peek().get(parameterIndex);
}

/**
* Exits the innermost correlation scope, returning the correlation ids to attach to the
* operator that owns it.
Expand Down
40 changes: 29 additions & 11 deletions isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -498,17 +498,22 @@ public RexNode visit(Expression.IfThen expr, Context context) throws RuntimeExce

@Override
public RexNode visit(Expression.Lambda expr, Context context) throws RuntimeException {
List<RelDataType> parameterTypes =
expr.parameters().fields().stream()
.map(parameter -> typeConverter.toCalcite(typeFactory, parameter))
.collect(Collectors.toList());
List<RexLambdaRef> parameters =
IntStream.range(0, expr.parameters().fields().size())
.mapToObj(
i ->
new RexLambdaRef(
i,
"p" + i,
typeConverter.toCalcite(typeFactory, expr.parameters().fields().get(i))))
IntStream.range(0, parameterTypes.size())
.mapToObj(i -> new RexLambdaRef(i, "p" + i, parameterTypes.get(i)))
.collect(Collectors.toList());

RexNode body = expr.body().accept(this, context);
context.enterLambdaScope(parameterTypes);
RexNode body;
try {
body = expr.body().accept(this, context);
} finally {
context.exitLambdaScope();
}

return rexBuilder.makeLambdaCall(body, parameters);
}
Expand Down Expand Up @@ -816,6 +821,10 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti
final FieldReference.StructField field = (FieldReference.StructField) segment;
rexInputRef =
new RexInputRef(field.offset(), typeConverter.toCalcite(typeFactory, expr.getType()));
observeType(
expr,
TypeObservation.Source.FIELD_REFERENCE,
() -> context.getInputFieldType(field.offset()));
} else {
throw new IllegalArgumentException("Unhandled type: " + segment);
}
Expand All @@ -834,8 +843,11 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti
final CorrelationId correlationId =
context.correlationIdForAnchor(
anchor, () -> relNodeConverter.getRelBuilder().getCluster().createCorrel());
return rexBuilder.makeFieldAccess(
rexBuilder.makeCorrel(rowType, correlationId), field.offset());
RexNode fieldAccess =
rexBuilder.makeFieldAccess(
rexBuilder.makeCorrel(rowType, correlationId), field.offset());
observeType(expr, TypeObservation.Source.FIELD_REFERENCE, fieldAccess::getType);
return fieldAccess;
} else {
throw new IllegalArgumentException("Unhandled type: " + segment);
}
Expand All @@ -852,7 +864,13 @@ public RexNode visit(FieldReference expr, Context context) throws RuntimeExcepti
if (segment instanceof FieldReference.StructField) {
final FieldReference.StructField field = (FieldReference.StructField) segment;
RelDataType calciteType = typeConverter.toCalcite(typeFactory, expr.getType());
return new RexLambdaRef(field.offset(), "p" + field.offset(), calciteType);
RexLambdaRef lambdaRef =
new RexLambdaRef(field.offset(), "p" + field.offset(), calciteType);
observeType(
expr,
TypeObservation.Source.FIELD_REFERENCE,
() -> context.getLambdaParameterType(field.offset()));
return lambdaRef;
} else {
throw new IllegalArgumentException("Unhandled type: " + segment);
}
Expand Down
5 changes: 4 additions & 1 deletion isthmus/src/main/java/io/substrait/isthmus/expression/TypeObservation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,10 @@ public enum Source {
SCALAR_FUNCTION,

/** A window function invocation. */
WINDOW_FUNCTION
WINDOW_FUNCTION,

/** A field reference. */
FIELD_REFERENCE
}

private final Source source;
Expand Down
Loading
Loading