diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 141ecf473..b7b010045 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -2013,6 +2013,46 @@ public Expression.WindowFunctionInvocation windowFn( WindowBound lowerBound, WindowBound upperBound, Expression... args) { + return windowFn( + urn, + key, + outputType, + aggregationPhase, + invocation, + Collections.emptyList(), + boundsType, + lowerBound, + upperBound, + args); + } + + /** + * Creates a window function invocation with specified arguments, window bounds, and ordering. + * + * @param urn the URN of the extension containing the function + * @param key the function key (name and signature) + * @param outputType the output type of the function + * @param aggregationPhase the aggregation phase + * @param invocation the aggregation invocation mode + * @param sort the ordering expressions for the window, required by a RANGE bound with a Preceding + * or Following side + * @param boundsType the type of window bounds + * @param lowerBound the lower bound of the window + * @param upperBound the upper bound of the window + * @param args the arguments to pass to the function + * @return a new {@link Expression.WindowFunctionInvocation} + */ + public Expression.WindowFunctionInvocation windowFn( + String urn, + String key, + Type outputType, + Expression.AggregationPhase aggregationPhase, + Expression.AggregationInvocation invocation, + List sort, + Expression.WindowBoundsType boundsType, + WindowBound lowerBound, + WindowBound upperBound, + Expression... args) { SimpleExtension.WindowFunctionVariant declaration = extensions.getWindowFunction(SimpleExtension.FunctionAnchor.of(urn, key)); return Expression.WindowFunctionInvocation.builder() @@ -2020,6 +2060,7 @@ public Expression.WindowFunctionInvocation windowFn( .outputType(outputType) .aggregationPhase(aggregationPhase) .invocation(invocation) + .sort(sort) .boundsType(boundsType) .lowerBound(lowerBound) .upperBound(upperBound) diff --git a/core/src/main/java/io/substrait/expression/Expression.java b/core/src/main/java/io/substrait/expression/Expression.java index 08a3f5cc0..88489f75c 100644 --- a/core/src/main/java/io/substrait/expression/Expression.java +++ b/core/src/main/java/io/substrait/expression/Expression.java @@ -1629,8 +1629,9 @@ public Type getType() { public abstract AggregationInvocation invocation(); /** - * Validates that variadic arguments satisfy the parameter consistency requirement, and that - * {@code bounds_type} is set whenever a window bound requires it. + * Validates that variadic arguments satisfy the parameter consistency requirement, that {@code + * bounds_type} is set whenever a window bound requires it, and that a RANGE bound with a + * Preceding or Following side has exactly one, non-CLUSTERED ordering expression. * *

When CONSISTENT, all variadic arguments must have the same type (ignoring nullability). * When INCONSISTENT, arguments can have different types. @@ -1639,6 +1640,7 @@ public Type getType() { protected void check() { VariadicParameterConsistencyValidator.validate(declaration(), arguments()); WindowBound.checkBoundsType(boundsType(), lowerBound(), upperBound()); + WindowBound.checkRangeOrdering(boundsType(), lowerBound(), upperBound(), sort()); } /** diff --git a/core/src/main/java/io/substrait/expression/WindowBound.java b/core/src/main/java/io/substrait/expression/WindowBound.java index fc03f0e75..1294e8d46 100644 --- a/core/src/main/java/io/substrait/expression/WindowBound.java +++ b/core/src/main/java/io/substrait/expression/WindowBound.java @@ -1,5 +1,6 @@ package io.substrait.expression; +import java.util.List; import java.util.Optional; import org.immutables.value.Value; @@ -62,6 +63,46 @@ static void checkBoundsType( } } + /** + * Validates a RANGE window's ordering against its bounds, per the spec's rule that a RANGE frame + * with a {@link Preceding} or {@link Following} bound must have exactly one ordering expression, + * which must not use {@code SORT_DIRECTION_CLUSTERED}. + * + * @param boundsType the window's bounds type + * @param lowerBound the window's lower bound + * @param upperBound the window's upper bound + * @param sorts the window's ordering expressions + * @throws IllegalArgumentException if {@code boundsType} is {@code RANGE}, either bound is {@link + * Preceding} or {@link Following}, and {@code sorts} does not contain exactly one ordering + * expression, or that expression uses {@code SORT_DIRECTION_CLUSTERED} + */ + static void checkRangeOrdering( + Expression.WindowBoundsType boundsType, + WindowBound lowerBound, + WindowBound upperBound, + List sorts) { + boolean needsSingleOrdering = + boundsType == Expression.WindowBoundsType.RANGE + && (lowerBound instanceof Preceding + || lowerBound instanceof Following + || upperBound instanceof Preceding + || upperBound instanceof Following); + if (!needsSingleOrdering) { + return; + } + if (sorts.size() != 1) { + throw new IllegalArgumentException( + "a RANGE bound with a Preceding or Following side requires exactly one ordering" + + " expression, but found " + + sorts.size()); + } + if (sorts.get(0).direction() == Expression.SortDirection.CLUSTERED) { + throw new IllegalArgumentException( + "a RANGE bound with a Preceding or Following side cannot use" + + " SORT_DIRECTION_CLUSTERED for its ordering expression"); + } + } + /** * Visitor over the concrete {@link WindowBound} kinds. * diff --git a/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java b/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java index 1338190a2..d8121de50 100644 --- a/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java +++ b/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java @@ -43,6 +43,21 @@ public abstract class ConsistentPartitionWindow extends SingleInputRel implement */ public abstract List getSorts(); + /** + * Validates that a RANGE bound with a Preceding or Following side has exactly one, non-CLUSTERED + * ordering expression, for every window function invocation. + */ + @Value.Check + protected void check() { + for (WindowRelFunctionInvocation windowFunction : getWindowFunctions()) { + WindowBound.checkRangeOrdering( + windowFunction.boundsType(), + windowFunction.lowerBound(), + windowFunction.upperBound(), + getSorts()); + } + } + /** * Derives the output record type by appending window outputs to the input type. * diff --git a/core/src/test/java/io/substrait/relation/ExpressionCopyOnWriteVisitorTest.java b/core/src/test/java/io/substrait/relation/ExpressionCopyOnWriteVisitorTest.java index 580cc896f..e683402db 100644 --- a/core/src/test/java/io/substrait/relation/ExpressionCopyOnWriteVisitorTest.java +++ b/core/src/test/java/io/substrait/relation/ExpressionCopyOnWriteVisitorTest.java @@ -98,7 +98,12 @@ void windowFunctionBoundOffsetsAreRewritten() { .declaration(declaration) .arguments(Collections.emptyList()) .partitionBy(Collections.emptyList()) - .sort(Collections.emptyList()) + .sort( + Collections.singletonList( + Expression.SortField.builder() + .expr(sb.i32(1)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build())) .outputType(R.I64) .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) .invocation(Expression.AggregationInvocation.ALL) @@ -114,6 +119,12 @@ void windowFunctionBoundOffsetsAreRewritten() { Optional.of( Expression.WindowFunctionInvocation.builder() .from(wfi) + .sort( + Collections.singletonList( + Expression.SortField.builder() + .expr(sb.i32(-1)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build())) .lowerBound(WindowBound.Preceding.of(sb.i32(-5))) .upperBound(WindowBound.Following.of(sb.i32(-7))) .build()), diff --git a/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java b/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java index 3c70bee82..61a6203c1 100644 --- a/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java +++ b/core/src/test/java/io/substrait/relation/OuterReferenceConverterTest.java @@ -308,7 +308,13 @@ void outerReferenceInsideWindowBoundOffsetIsConverted() { .declaration(declaration) .arguments(List.of(sb.fieldReference(input2, 0))) .partitionBy(Collections.emptyList()) - .sort(Collections.emptyList()) + .sort( + List.of( + Expression.SortField.builder() + .expr(sb.fieldReference(input2, 0)) + .direction( + Expression.SortDirection.ASC_NULLS_FIRST) + .build())) .outputType(TypeCreator.NULLABLE.I64) .aggregationPhase( Expression.AggregationPhase.INITIAL_TO_RESULT) diff --git a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java index 7856a81d1..7dc79529c 100644 --- a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java +++ b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java @@ -41,7 +41,7 @@ void consistentPartitionWindowBoundOffsetsAreRewritten() { extensions.getWindowFunction( SimpleExtension.FunctionAnchor.of( DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); - Rel input = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I64)); + Rel input = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I32)); ConsistentPartitionWindow window = ConsistentPartitionWindow.builder() .input(input) @@ -57,6 +57,12 @@ void consistentPartitionWindowBoundOffsetsAreRewritten() { .upperBound(WindowBound.Following.of(sb.i32(7))) .boundsType(Expression.WindowBoundsType.RANGE) .build())) + .sorts( + Arrays.asList( + Expression.SortField.builder() + .expr(sb.fieldReference(input, 0)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build())) .build(); Optional rewritten = diff --git a/core/src/test/java/io/substrait/type/proto/ConsistentPartitionWindowRelRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/ConsistentPartitionWindowRelRoundtripTest.java index bbaafa66d..e1d10bdd9 100644 --- a/core/src/test/java/io/substrait/type/proto/ConsistentPartitionWindowRelRoundtripTest.java +++ b/core/src/test/java/io/substrait/type/proto/ConsistentPartitionWindowRelRoundtripTest.java @@ -202,12 +202,19 @@ void windowFunctionInvocationRoundtripWithNonLiteralOffsetExpr() { DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); // Unlike the relation-level fixture above, this bare expression has no enclosing relation to // resolve field references against, so the offset is a scalar function call over literals. + // A RANGE bound with a Preceding side requires exactly one ordering expression, carried here + // directly on the invocation since there is no enclosing relation to hold it. Expression.WindowFunctionInvocation wfi = Expression.WindowFunctionInvocation.builder() .declaration(windowFunctionDeclaration) .arguments(Collections.emptyList()) .partitionBy(Collections.emptyList()) - .sort(Collections.emptyList()) + .sort( + Collections.singletonList( + Expression.SortField.builder() + .expr(sb.i64(1)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build())) .outputType(R.I64) .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) .invocation(Expression.AggregationInvocation.ALL) @@ -254,6 +261,108 @@ void boundsTypeUnspecifiedWithRealBoundIsRejected() { assertThrows(IllegalArgumentException.class, invocationBuilder::build); } + @Test + void rangePrecedingWithoutASingleOrderingExpressionIsRejected() { + SimpleExtension.WindowFunctionVariant windowFunctionDeclaration = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + Rel input = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I64)); + // A RANGE bound with a Preceding side requires exactly one ordering expression on the + // enclosing relation; this fixture has none. The check runs in a @Value.Check, so it fires at + // construction time rather than only when a plan is later read back from proto. + ImmutableConsistentPartitionWindow.Builder relBuilder = + ConsistentPartitionWindow.builder() + .input(input) + .windowFunctions( + Arrays.asList( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(windowFunctionDeclaration) + .arguments(Arrays.asList(sb.fieldReference(input, 0))) + .outputType(R.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .lowerBound(WindowBound.Preceding.of(5)) + .upperBound(WindowBound.CURRENT_ROW) + .boundsType(Expression.WindowBoundsType.RANGE) + .build())); + + assertThrows(IllegalArgumentException.class, relBuilder::build); + } + + @Test + void rangePrecedingWithTwoOrderingExpressionsIsRejected() { + SimpleExtension.WindowFunctionVariant windowFunctionDeclaration = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + Rel input = + sb.namedScan(Arrays.asList("test"), Arrays.asList("a", "b"), Arrays.asList(R.I64, R.I64)); + // A RANGE bound with a Preceding side requires exactly one ordering expression; this fixture + // has two. + ImmutableConsistentPartitionWindow.Builder relBuilder = + ConsistentPartitionWindow.builder() + .input(input) + .windowFunctions( + Arrays.asList( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(windowFunctionDeclaration) + .arguments(Arrays.asList(sb.fieldReference(input, 0))) + .outputType(R.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .lowerBound(WindowBound.Preceding.of(5)) + .upperBound(WindowBound.CURRENT_ROW) + .boundsType(Expression.WindowBoundsType.RANGE) + .build())) + .sorts( + Arrays.asList( + Expression.SortField.builder() + .expr(sb.fieldReference(input, 0)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build(), + Expression.SortField.builder() + .expr(sb.fieldReference(input, 1)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build())); + + assertThrows(IllegalArgumentException.class, relBuilder::build); + } + + @Test + void rangePrecedingWithClusteredOrderingIsRejected() { + SimpleExtension.WindowFunctionVariant windowFunctionDeclaration = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + Rel input = sb.namedScan(Arrays.asList("test"), Arrays.asList("a"), Arrays.asList(R.I64)); + // A RANGE bound with a Preceding side cannot use SORT_DIRECTION_CLUSTERED for its ordering + // expression. + ImmutableConsistentPartitionWindow.Builder relBuilder = + ConsistentPartitionWindow.builder() + .input(input) + .windowFunctions( + Arrays.asList( + ConsistentPartitionWindow.WindowRelFunctionInvocation.builder() + .declaration(windowFunctionDeclaration) + .arguments(Arrays.asList(sb.fieldReference(input, 0))) + .outputType(R.I64) + .aggregationPhase(Expression.AggregationPhase.INITIAL_TO_RESULT) + .invocation(Expression.AggregationInvocation.ALL) + .lowerBound(WindowBound.Preceding.of(5)) + .upperBound(WindowBound.CURRENT_ROW) + .boundsType(Expression.WindowBoundsType.RANGE) + .build())) + .sorts( + Arrays.asList( + Expression.SortField.builder() + .expr(sb.fieldReference(input, 0)) + .direction(Expression.SortDirection.CLUSTERED) + .build())); + + assertThrows(IllegalArgumentException.class, relBuilder::build); + } + @Test void boundsTypeUnspecifiedWithUnboundedBoundsIsAccepted() { SimpleExtension.WindowFunctionVariant windowFunctionDeclaration =