diff --git a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java index 141ecf473..c18954bd9 100644 --- a/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java +++ b/core/src/main/java/io/substrait/dsl/SubstraitBuilder.java @@ -1990,7 +1990,9 @@ public Expression.ScalarFunctionInvocation scalarFn( } /** - * Creates a window function invocation with specified arguments and window bounds. + * Creates a window function invocation with specified arguments and window bounds. Supplies no + * ordering expressions, so a RANGE bound with a Preceding or Following side is rejected outright; + * use the {@code sort}-carrying overload for that. * * @param urn the URN of the extension containing the function * @param key the function key (name and signature) @@ -2013,6 +2015,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 +2062,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..767003285 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,8 @@ public Type getType() { protected void check() { VariadicParameterConsistencyValidator.validate(declaration(), arguments()); WindowBound.checkBoundsType(boundsType(), lowerBound(), upperBound()); + WindowBound.checkRangeOrdering( + boundsType(), lowerBound(), upperBound(), sort(), declaration().key()); } /** diff --git a/core/src/main/java/io/substrait/expression/WindowBound.java b/core/src/main/java/io/substrait/expression/WindowBound.java index fc03f0e75..6ec2e994f 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,50 @@ 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 + * @param function identifies the window function being validated, for the exception message + * @throws IllegalArgumentException if {@code boundsType} is {@code RANGE} and either bound is + * {@link Preceding} or {@link Following}, and {@code sorts} does not hold exactly one + * ordering expression whose direction is not {@code SORT_DIRECTION_CLUSTERED} + */ + static void checkRangeOrdering( + Expression.WindowBoundsType boundsType, + WindowBound lowerBound, + WindowBound upperBound, + List sorts, + String function) { + 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( + function + + ": 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( + function + + ": 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. * @@ -123,9 +168,8 @@ abstract class Preceding implements WindowBound { public abstract Expression offset(); /** - * Creates a {@link Preceding} bound from a literal row offset. For {@code BOUNDS_TYPE_ROWS} - * only: a RANGE bound's offset must be type-compatible with the ordering expression, so use - * {@link #of(Expression)} there. + * Creates a {@link Preceding} bound from a literal {@code i64} row offset. Valid for ROWS, or + * for RANGE over an {@code i64} ordering expression; use {@link #of(Expression)} otherwise. * * @param offset the row offset preceding the current row * @return the preceding bound @@ -161,9 +205,8 @@ abstract class Following implements WindowBound { public abstract Expression offset(); /** - * Creates a {@link Following} bound from a literal row offset. For {@code BOUNDS_TYPE_ROWS} - * only: a RANGE bound's offset must be type-compatible with the ordering expression, so use - * {@link #of(Expression)} there. + * Creates a {@link Following} bound from a literal {@code i64} row offset. Valid for ROWS, or + * for RANGE over an {@code i64} ordering expression; use {@link #of(Expression)} otherwise. * * @param offset the row offset following the current row * @return the following bound diff --git a/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java b/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java index 1338190a2..98ca0c2e2 100644 --- a/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java +++ b/core/src/main/java/io/substrait/relation/ConsistentPartitionWindow.java @@ -43,6 +43,24 @@ 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() { + List windowFunctions = getWindowFunctions(); + for (int i = 0; i < windowFunctions.size(); i++) { + WindowRelFunctionInvocation windowFunction = windowFunctions.get(i); + WindowBound.checkRangeOrdering( + windowFunction.boundsType(), + windowFunction.lowerBound(), + windowFunction.upperBound(), + getSorts(), + "window function " + i + " (" + windowFunction.declaration().key() + ")"); + } + } + /** * Derives the output record type by appending window outputs to the input type. * diff --git a/core/src/test/java/io/substrait/dsl/SubstraitBuilderTest.java b/core/src/test/java/io/substrait/dsl/SubstraitBuilderTest.java index ef0b3afc1..3dbb512fb 100644 --- a/core/src/test/java/io/substrait/dsl/SubstraitBuilderTest.java +++ b/core/src/test/java/io/substrait/dsl/SubstraitBuilderTest.java @@ -3,11 +3,13 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import io.substrait.TestBase; import io.substrait.expression.AggregateFunctionInvocation; import io.substrait.expression.Expression; import io.substrait.expression.FieldReference; +import io.substrait.expression.WindowBound; import io.substrait.extension.DefaultExtensionCatalog; import io.substrait.extension.SimpleExtension; import io.substrait.plan.Plan; @@ -227,6 +229,43 @@ void testBooleanLogic() { assertNotNull(builder.not(b1)); assertNotNull(builder.isNull(b1)); } + + @Test + void testWindowFunctionWithOrdering() { + // The single shape the sorts-carrying overload exists for: a RANGE bound with a Preceding + // side, which requires exactly one ordering expression. + final NamedScan scan = createSimpleScan(); + final Expression.WindowFunctionInvocation windowFn = + builder.windowFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "lead:any", + Type.I32.builder().nullable(false).build(), + Expression.AggregationPhase.INITIAL_TO_RESULT, + Expression.AggregationInvocation.ALL, + builder.sortFields(scan, 0), + Expression.WindowBoundsType.RANGE, + WindowBound.Preceding.of(builder.i32(5)), + WindowBound.CURRENT_ROW, + builder.fieldReference(scan, 0)); + + assertNotNull(windowFn); + assertEquals(1, windowFn.sort().size()); + + // The no-sorts overload cannot express this shape at all. + assertThrows( + IllegalArgumentException.class, + () -> + builder.windowFn( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, + "lead:any", + Type.I32.builder().nullable(false).build(), + Expression.AggregationPhase.INITIAL_TO_RESULT, + Expression.AggregationInvocation.ALL, + Expression.WindowBoundsType.RANGE, + WindowBound.Preceding.of(builder.i32(5)), + WindowBound.CURRENT_ROW, + builder.fieldReference(scan, 0))); + } } @Nested 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 fd3c779ea..64662265f 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 f4a56b062..f47674904 100644 --- a/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java +++ b/core/src/test/java/io/substrait/relation/RelCopyOnWriteVisitorTest.java @@ -61,12 +61,13 @@ private ConsistentPartitionWindow windowOver(Rel input, WindowBound lower, Windo .upperBound(upper) .boundsType(Expression.WindowBoundsType.RANGE) .build())) + .sorts(sb.sortFields(input, 0)) .build(); } @Test void consistentPartitionWindowBoundOffsetsAreRewritten() { - 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 = windowOver(input, WindowBound.Preceding.of(sb.i32(5)), WindowBound.Following.of(sb.i32(7))); 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..7b2c0c71e 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,128 @@ 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(sb.sortFields(input, 0, 1)); + + 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 rangePrecedingWithTwoOrderingExpressionsOnAnInvocationIsRejected() { + SimpleExtension.WindowFunctionVariant declaration = + extensions.getWindowFunction( + SimpleExtension.FunctionAnchor.of( + DefaultExtensionCatalog.FUNCTIONS_ARITHMETIC, "lead:any")); + Expression.SortField sort = + Expression.SortField.builder() + .expr(sb.i64(1)) + .direction(Expression.SortDirection.ASC_NULLS_FIRST) + .build(); + + assertThrows( + IllegalArgumentException.class, + () -> + Expression.WindowFunctionInvocation.builder() + .declaration(declaration) + .arguments(Collections.emptyList()) + .partitionBy(Collections.emptyList()) + .sort(Arrays.asList(sort, sort)) + .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()); + } + @Test void boundsTypeUnspecifiedWithUnboundedBoundsIsAccepted() { SimpleExtension.WindowFunctionVariant windowFunctionDeclaration =