feat(isthmus): add WindowRelFunctionConverter #234
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,32 @@ | ||
| package io.substrait.isthmus.expression; | ||
|
|
||
| import io.substrait.expression.Expression; | ||
| import org.apache.calcite.rel.RelFieldCollation; | ||
| import org.apache.calcite.rex.RexFieldCollation; | ||
|
|
||
| public class SortFieldConverter { | ||
|
|
||
| /** Converts a {@link RexFieldCollation} to a {@link Expression.SortField}. */ | ||
| public static Expression.SortField toSortField( | ||
| RexFieldCollation rexFieldCollation, RexExpressionConverter rexExpressionConverter) { | ||
| var expr = rexFieldCollation.left.accept(rexExpressionConverter); | ||
| var rexDirection = rexFieldCollation.getDirection(); | ||
| Expression.SortDirection direction = | ||
| switch (rexDirection) { | ||
| case ASCENDING -> rexFieldCollation.getNullDirection() | ||
| == RelFieldCollation.NullDirection.LAST | ||
| ? Expression.SortDirection.ASC_NULLS_LAST | ||
| : Expression.SortDirection.ASC_NULLS_FIRST; | ||
| case DESCENDING -> rexFieldCollation.getNullDirection() | ||
| == RelFieldCollation.NullDirection.LAST | ||
| ? Expression.SortDirection.DESC_NULLS_LAST | ||
| : Expression.SortDirection.DESC_NULLS_FIRST; | ||
| default -> throw new IllegalArgumentException( | ||
| String.format( | ||
| "Unexpected RelFieldCollation.Direction:%s enum at the RexFieldCollation!", | ||
| rexDirection)); | ||
| }; | ||
|
|
||
| return Expression.SortField.builder().expr(expr).direction(direction).build(); | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| package io.substrait.isthmus.expression; | ||
|
|
||
| import io.substrait.expression.WindowBound; | ||
| import java.math.BigDecimal; | ||
| import org.apache.calcite.rex.RexLiteral; | ||
| import org.apache.calcite.rex.RexWindowBound; | ||
| import org.apache.calcite.sql.type.SqlTypeName; | ||
|
|
||
| public class WindowBoundConverter { | ||
|
|
||
| /** Converts a {@link RexWindowBound} to a {@link WindowBound}. */ | ||
| public static WindowBound toWindowBound(RexWindowBound rexWindowBound) { | ||
|
There was a problem hiding this comment. ✨ Good call extracting this. It should be useful to any downstream uses who wish to implement their own conversion. |
||
| if (rexWindowBound.isCurrentRow()) { | ||
| return WindowBound.CURRENT_ROW; | ||
| } | ||
| if (rexWindowBound.isUnbounded()) { | ||
| return WindowBound.UNBOUNDED; | ||
| } else { | ||
| if (rexWindowBound.getOffset() instanceof RexLiteral literal | ||
| && SqlTypeName.EXACT_TYPES.contains(literal.getTypeName())) { | ||
| BigDecimal offset = (BigDecimal) literal.getValue4(); | ||
| if (rexWindowBound.isPreceding()) { | ||
| return WindowBound.Preceding.of(offset.longValue()); | ||
| } | ||
| if (rexWindowBound.isFollowing()) { | ||
| return WindowBound.Following.of(offset.longValue()); | ||
| } | ||
| throw new IllegalStateException( | ||
| "window bound was none of CURRENT ROW, UNBOUNDED, PRECEDING or FOLLOWING"); | ||
| } | ||
| throw new IllegalArgumentException( | ||
| String.format( | ||
| "substrait only supports integer window offsets. Received: %s", | ||
| rexWindowBound.getOffset().getKind())); | ||
| } | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,136 @@ | ||
| package io.substrait.isthmus.expression; | ||
|
|
||
| import static io.substrait.isthmus.expression.WindowBoundConverter.toWindowBound; | ||
|
|
||
| import com.google.common.collect.ImmutableList; | ||
| import io.substrait.expression.Expression; | ||
| import io.substrait.expression.ExpressionCreator; | ||
| import io.substrait.expression.FunctionArg; | ||
| import io.substrait.expression.WindowBound; | ||
| import io.substrait.extension.SimpleExtension; | ||
| import io.substrait.isthmus.AggregateFunctions; | ||
| import io.substrait.relation.ConsistentPartitionWindow; | ||
| import io.substrait.type.Type; | ||
| import java.util.List; | ||
| import java.util.Optional; | ||
| import java.util.function.Function; | ||
| import java.util.stream.Stream; | ||
| import org.apache.calcite.rel.core.Window; | ||
| import org.apache.calcite.rel.type.RelDataType; | ||
| import org.apache.calcite.rel.type.RelDataTypeFactory; | ||
| import org.apache.calcite.rex.RexNode; | ||
| import org.apache.calcite.rex.RexWindowBound; | ||
| import org.apache.calcite.sql.SqlAggFunction; | ||
|
|
||
| public class WindowRelFunctionConverter | ||
| extends FunctionConverter< | ||
| SimpleExtension.WindowFunctionVariant, | ||
| ConsistentPartitionWindow.WindowRelFunctionInvocation, | ||
| WindowRelFunctionConverter.WrappedWindowRelCall> { | ||
|
|
||
| @Override | ||
| protected ImmutableList<FunctionMappings.Sig> getSigs() { | ||
| return FunctionMappings.WINDOW_SIGS; | ||
| } | ||
|
|
||
| public WindowRelFunctionConverter( | ||
| List<SimpleExtension.WindowFunctionVariant> functions, RelDataTypeFactory typeFactory) { | ||
| super(functions, typeFactory); | ||
| } | ||
|
|
||
| @Override | ||
| protected ConsistentPartitionWindow.WindowRelFunctionInvocation generateBinding( | ||
| WrappedWindowRelCall call, | ||
| SimpleExtension.WindowFunctionVariant function, | ||
| List<FunctionArg> arguments, | ||
| Type outputType) { | ||
| Window.RexWinAggCall over = call.getWinAggCall(); | ||
|
|
||
| Expression.AggregationInvocation invocation = | ||
| over.distinct | ||
| ? Expression.AggregationInvocation.DISTINCT | ||
| : Expression.AggregationInvocation.ALL; | ||
|
|
||
| // Calcite only supports ROW or RANGE mode | ||
| Expression.WindowBoundsType boundsType = | ||
| call.isRows() ? Expression.WindowBoundsType.ROWS : Expression.WindowBoundsType.RANGE; | ||
| WindowBound lowerBound = toWindowBound(call.getLowerBound()); | ||
| WindowBound upperBound = toWindowBound(call.getUpperBound()); | ||
|
|
||
| return ExpressionCreator.windowRelFunction( | ||
| function, | ||
| outputType, | ||
| Expression.AggregationPhase.INITIAL_TO_RESULT, | ||
| invocation, | ||
| boundsType, | ||
| lowerBound, | ||
| upperBound, | ||
| arguments); | ||
| } | ||
|
|
||
| public Optional<ConsistentPartitionWindow.WindowRelFunctionInvocation> convert( | ||
| Window.RexWinAggCall winAggCall, | ||
| RexWindowBound lowerBound, | ||
| RexWindowBound upperBound, | ||
| boolean isRows, | ||
| Function<RexNode, Expression> topLevelConverter) { | ||
| var aggFunction = (SqlAggFunction) winAggCall.getOperator(); | ||
|
|
||
| SqlAggFunction lookupFunction = | ||
| AggregateFunctions.toSubstraitAggVariant(aggFunction).orElse(aggFunction); | ||
| FunctionFinder m = signatures.get(lookupFunction); | ||
| if (m == null) { | ||
| return Optional.empty(); | ||
| } | ||
| if (!m.allowedArgCount(winAggCall.getOperands().size())) { | ||
| return Optional.empty(); | ||
| } | ||
|
|
||
| var wrapped = new WrappedWindowRelCall(winAggCall, lowerBound, upperBound, isRows); | ||
| return m.attemptMatch(wrapped, topLevelConverter); | ||
| } | ||
|
|
||
| static class WrappedWindowRelCall implements GenericCall { | ||
| private final Window.RexWinAggCall winAggCall; | ||
| private final RexWindowBound lowerBound; | ||
| private final RexWindowBound upperBound; | ||
| private final boolean isRows; | ||
|
|
||
| private WrappedWindowRelCall( | ||
| Window.RexWinAggCall winAggCall, | ||
| RexWindowBound lowerBound, | ||
| RexWindowBound upperBound, | ||
| boolean isRows) { | ||
| this.winAggCall = winAggCall; | ||
| this.lowerBound = lowerBound; | ||
| this.upperBound = upperBound; | ||
| this.isRows = isRows; | ||
| } | ||
|
|
||
| @Override | ||
| public Stream<RexNode> getOperands() { | ||
| return winAggCall.getOperands().stream(); | ||
| } | ||
|
|
||
| @Override | ||
| public RelDataType getType() { | ||
| return winAggCall.getType(); | ||
| } | ||
|
|
||
| public Window.RexWinAggCall getWinAggCall() { | ||
| return winAggCall; | ||
| } | ||
|
|
||
| public RexWindowBound getLowerBound() { | ||
| return lowerBound; | ||
| } | ||
|
|
||
| public RexWindowBound getUpperBound() { | ||
| return upperBound; | ||
| } | ||
|
|
||
| public boolean isRows() { | ||
| return isRows; | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
✨ Good call extracting this. It should be useful to any downstream uses who wish to implement their own conversion.