Skip to content

Commit

Permalink
Support row pattern measures over window: grammmar and AST
Browse files Browse the repository at this point in the history
  • Loading branch information
@kasiafi
kasiafi committed Jun 21, 2021
1 parent dc84283 commit f98fe48
Show file tree
Hide file tree
Showing 11 changed files with 310 additions and 80 deletions.
7 changes: 7 additions & 0 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/ExpressionAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,7 @@
import io.trino.sql.tree.TryExpression;
import io.trino.sql.tree.WhenClause;
import io.trino.sql.tree.WindowFrame;
import io.trino.sql.tree.WindowOperation;
import io.trino.type.FunctionType;
import io.trino.type.TypeCoercion;
import io.trino.type.UnknownType;
Expand Down Expand Up @@ -1365,6 +1366,12 @@ private void analyzeFrameRangeOffset(Expression offsetValue, FrameBound.Type bou
frameBoundCalculations.put(NodeRef.of(offsetValue), function);
}

@Override
protected Type visitWindowOperation(WindowOperation node, StackableAstVisitorContext<Context> context)
{
throw semanticException(NOT_SUPPORTED, node, "Row pattern measures over window not yet supported");
}

public List<TypeSignatureProvider> getCallArgumentTypes(List<Expression> arguments, StackableAstVisitorContext<Context> context)
{
ImmutableList.Builder<TypeSignatureProvider> argumentTypesBuilder = ImmutableList.builder();
Expand Down
14 changes: 14 additions & 0 deletions core/trino-main/src/test/java/io/trino/sql/analyzer/TestAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1384,6 +1384,20 @@ public void testWindowFrameWithPatternRecognition()
.hasMessage("line 1:200: Pattern recognition in window not yet supported");
}

@Test
public void testMeasureOverWindow()
{
// in-line window specification
assertFails("SELECT last_z OVER () FROM (VALUES 1) t(z) ")
.hasErrorCode(NOT_SUPPORTED)
.hasMessage("line 1:8: Row pattern measures over window not yet supported");

// named window reference
assertFails("SELECT last_z OVER w FROM (VALUES 1) t(z) WINDOW w AS ()")
.hasErrorCode(NOT_SUPPORTED)
.hasMessage("line 1:8: Row pattern measures over window not yet supported");
}

@Test
public void testDistinctInWindowFunctionParameter()
{
Expand Down
1 change: 1 addition & 0 deletions core/trino-parser/src/main/antlr4/io/trino/sql/parser/SqlBase.g4
View file
Original file line number Diff line number Diff line change
Expand Up @@ -410,6 +410,7 @@ primaryExpression
| qualifiedName '(' ASTERISK ')' filter? over? #functionCall
| processingMode? qualifiedName '(' (setQuantifier? expression (',' expression)*)?
(ORDER BY sortItem (',' sortItem)*)? ')' filter? (nullTreatment? over)? #functionCall
| identifier over #measure
| identifier '->' expression #lambda
| '(' (identifier (',' identifier)*)? ')' '->' expression #lambda
| '(' query ')' #subqueryExpression
Expand Down
7 changes: 7 additions & 0 deletions core/trino-parser/src/main/java/io/trino/sql/ExpressionFormatter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@
import io.trino.sql.tree.WhenClause;
import io.trino.sql.tree.Window;
import io.trino.sql.tree.WindowFrame;
import io.trino.sql.tree.WindowOperation;
import io.trino.sql.tree.WindowReference;
import io.trino.sql.tree.WindowSpecification;

Expand Down Expand Up @@ -425,6 +426,12 @@ protected String visitFunctionCall(FunctionCall node, Void context)
return builder.toString();
}

@Override
protected String visitWindowOperation(WindowOperation node, Void context)
{
return process(node.getName(), context) + " OVER " + formatWindow(node.getWindow());
}

@Override
protected String visitLambdaExpression(LambdaExpression node, Void context)
{
Expand Down
7 changes: 7 additions & 0 deletions core/trino-parser/src/main/java/io/trino/sql/parser/AstBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@
import io.trino.sql.tree.Window;
import io.trino.sql.tree.WindowDefinition;
import io.trino.sql.tree.WindowFrame;
import io.trino.sql.tree.WindowOperation;
import io.trino.sql.tree.WindowReference;
import io.trino.sql.tree.WindowSpecification;
import io.trino.sql.tree.With;
Expand Down Expand Up @@ -2096,6 +2097,12 @@ else if (processingMode.FINAL() != null) {
visit(context.expression(), Expression.class));
}

@Override
public Node visitMeasure(SqlBaseParser.MeasureContext context)
{
return new WindowOperation(getLocation(context), (Identifier) visit(context.identifier()), (Window) visit(context.over()));
}

@Override
public Node visitLambda(SqlBaseParser.LambdaContext context)
{
Expand Down
5 changes: 5 additions & 0 deletions core/trino-parser/src/main/java/io/trino/sql/tree/AstVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,11 @@ protected R visitProcessingMode(ProcessingMode node, C context)
return visitNode(node, context);
}

protected R visitWindowOperation(WindowOperation node, C context)
{
return visitExpression(node, context);
}

protected R visitLambdaExpression(LambdaExpression node, C context)
{
return visitExpression(node, context);
Expand Down
9 changes: 9 additions & 0 deletions core/trino-parser/src/main/java/io/trino/sql/tree/DefaultTraversalVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,15 @@ protected Void visitFunctionCall(FunctionCall node, C context)
return null;
}

@Override
protected Void visitWindowOperation(WindowOperation node, C context)
{
process(node.getName(), context);
process((Node) node.getWindow(), context);

return null;
}

@Override
protected Void visitGroupingOperation(GroupingOperation node, C context)
{
Expand Down
5 changes: 5 additions & 0 deletions core/trino-parser/src/main/java/io/trino/sql/tree/ExpressionRewriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,11 @@ public Expression rewriteFunctionCall(FunctionCall node, C context, ExpressionTr
return rewriteExpression(node, context, treeRewriter);
}

public Expression rewriteWindowOperation(WindowOperation node, C context, ExpressionTreeRewriter<C> treeRewriter)
{
return rewriteExpression(node, context, treeRewriter);
}

public Expression rewriteLambdaExpression(LambdaExpression node, C context, ExpressionTreeRewriter<C> treeRewriter)
{
return rewriteExpression(node, context, treeRewriter);
Expand Down
187 changes: 107 additions & 80 deletions core/trino-parser/src/main/java/io/trino/sql/tree/ExpressionTreeRewriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -484,95 +484,22 @@ public Expression visitFunctionCall(FunctionCall node, Context<C> context)
filter = Optional.of(newFilterExpression);
}

Optional<Window> rewrittenWindow = node.getWindow();
if (node.getWindow().isPresent()) {
Window window = node.getWindow().get();

if (window instanceof WindowReference) {
WindowReference windowReference = (WindowReference) window;
Identifier rewrittenName = rewrite(windowReference.getName(), context.get());
if (windowReference.getName() != rewrittenName) {
rewrittenWindow = Optional.of(new WindowReference(rewrittenName));
}
}
else if (window instanceof WindowSpecification) {
WindowSpecification windowSpecification = (WindowSpecification) window;
Optional<Identifier> existingWindowName = windowSpecification.getExistingWindowName().map(name -> rewrite(name, context.get()));

List<Expression> partitionBy = rewrite(windowSpecification.getPartitionBy(), context);

Optional<OrderBy> orderBy = Optional.empty();
if (windowSpecification.getOrderBy().isPresent()) {
orderBy = Optional.of(rewriteOrderBy(windowSpecification.getOrderBy().get(), context));
}

Optional<WindowFrame> rewrittenFrame = windowSpecification.getFrame();
if (rewrittenFrame.isPresent()) {
WindowFrame frame = rewrittenFrame.get();

FrameBound start = frame.getStart();
if (start.getValue().isPresent()) {
Expression value = rewrite(start.getValue().get(), context.get());
if (value != start.getValue().get()) {
start = new FrameBound(start.getType(), value);
}
}

Optional<FrameBound> rewrittenEnd = frame.getEnd();
if (rewrittenEnd.isPresent()) {
Optional<Expression> value = rewrittenEnd.get().getValue();
if (value.isPresent()) {
Expression rewrittenValue = rewrite(value.get(), context.get());
if (rewrittenValue != value.get()) {
rewrittenEnd = Optional.of(new FrameBound(rewrittenEnd.get().getType(), rewrittenValue));
}
}
}

// Frame properties for row pattern matching are not rewritten. They are planned as parts of
// PatternRecognitionNode, and shouldn't be accessed past the Planner phase.
// There are nested expressions in Measures and VariableDefinitions. They are not rewritten by default.
// Rewriting them requires special handling of DereferenceExpression, aware of pattern labels.
if (!frame.getMeasures().isEmpty() ||
frame.getAfterMatchSkipTo().isPresent() ||
frame.getPatternSearchMode().isPresent() ||
frame.getPattern().isPresent() ||
!frame.getSubsets().isEmpty() ||
!frame.getVariableDefinitions().isEmpty()) {
throw new UnsupportedOperationException("Cannot rewrite pattern recognition clauses in window");
}

if ((frame.getStart() != start) || !sameElements(frame.getEnd(), rewrittenEnd)) {
rewrittenFrame = Optional.of(new WindowFrame(
frame.getType(),
start,
rewrittenEnd,
frame.getMeasures(),
frame.getAfterMatchSkipTo(),
frame.getPatternSearchMode(),
frame.getPattern(),
frame.getSubsets(),
frame.getVariableDefinitions()));
}
}

if (!sameElements(windowSpecification.getExistingWindowName(), existingWindowName) ||
!sameElements(windowSpecification.getPartitionBy(), partitionBy) ||
!sameElements(windowSpecification.getOrderBy(), orderBy) ||
!sameElements(windowSpecification.getFrame(), rewrittenFrame)) {
rewrittenWindow = Optional.of(new WindowSpecification(existingWindowName, partitionBy, orderBy, rewrittenFrame));
}
Optional<Window> window = node.getWindow();
if (window.isPresent()) {
Window rewrittenWindow = rewriteWindow(window.get(), context);
if (rewrittenWindow != window.get()) {
window = Optional.of(rewrittenWindow);
}
}

List<Expression> arguments = rewrite(node.getArguments(), context);

if (!sameElements(node.getArguments(), arguments) || !sameElements(rewrittenWindow, node.getWindow())
if (!sameElements(node.getArguments(), arguments) || !sameElements(window, node.getWindow())
|| !sameElements(filter, node.getFilter())) {
return new FunctionCall(
node.getLocation(),
node.getName(),
rewrittenWindow,
window,
filter,
node.getOrderBy().map(orderBy -> rewriteOrderBy(orderBy, context)),
node.isDistinct(),
Expand Down Expand Up @@ -609,6 +536,106 @@ private List<SortItem> rewriteSortItems(List<SortItem> sortItems, Context<C> con
return rewrittenSortItems.build();
}

private Window rewriteWindow(Window window, Context<C> context)
{
if (window instanceof WindowReference) {
WindowReference windowReference = (WindowReference) window;
Identifier rewrittenName = rewrite(windowReference.getName(), context.get());
if (windowReference.getName() != rewrittenName) {
return new WindowReference(rewrittenName);
}
return window;
}

WindowSpecification windowSpecification = (WindowSpecification) window;
Optional<Identifier> existingWindowName = windowSpecification.getExistingWindowName().map(name -> rewrite(name, context.get()));

List<Expression> partitionBy = rewrite(windowSpecification.getPartitionBy(), context);

Optional<OrderBy> orderBy = Optional.empty();
if (windowSpecification.getOrderBy().isPresent()) {
orderBy = Optional.of(rewriteOrderBy(windowSpecification.getOrderBy().get(), context));
}

Optional<WindowFrame> rewrittenFrame = windowSpecification.getFrame();
if (rewrittenFrame.isPresent()) {
WindowFrame frame = rewrittenFrame.get();

FrameBound start = frame.getStart();
if (start.getValue().isPresent()) {
Expression value = rewrite(start.getValue().get(), context.get());
if (value != start.getValue().get()) {
start = new FrameBound(start.getType(), value);
}
}

Optional<FrameBound> rewrittenEnd = frame.getEnd();
if (rewrittenEnd.isPresent()) {
Optional<Expression> value = rewrittenEnd.get().getValue();
if (value.isPresent()) {
Expression rewrittenValue = rewrite(value.get(), context.get());
if (rewrittenValue != value.get()) {
rewrittenEnd = Optional.of(new FrameBound(rewrittenEnd.get().getType(), rewrittenValue));
}
}
}

// Frame properties for row pattern matching are not rewritten. They are planned as parts of
// PatternRecognitionNode, and shouldn't be accessed past the Planner phase.
// There are nested expressions in Measures and VariableDefinitions. They are not rewritten by default.
// Rewriting them requires special handling of DereferenceExpression, aware of pattern labels.
if (!frame.getMeasures().isEmpty() ||
frame.getAfterMatchSkipTo().isPresent() ||
frame.getPatternSearchMode().isPresent() ||
frame.getPattern().isPresent() ||
!frame.getSubsets().isEmpty() ||
!frame.getVariableDefinitions().isEmpty()) {
throw new UnsupportedOperationException("cannot rewrite pattern recognition clauses in window");
}

if ((frame.getStart() != start) || !sameElements(frame.getEnd(), rewrittenEnd)) {
rewrittenFrame = Optional.of(new WindowFrame(
frame.getType(),
start,
rewrittenEnd,
frame.getMeasures(),
frame.getAfterMatchSkipTo(),
frame.getPatternSearchMode(),
frame.getPattern(),
frame.getSubsets(),
frame.getVariableDefinitions()));
}
}

if (!sameElements(windowSpecification.getExistingWindowName(), existingWindowName) ||
!sameElements(windowSpecification.getPartitionBy(), partitionBy) ||
!sameElements(windowSpecification.getOrderBy(), orderBy) ||
!sameElements(windowSpecification.getFrame(), rewrittenFrame)) {
return new WindowSpecification(existingWindowName, partitionBy, orderBy, rewrittenFrame);
}
return window;
}

@Override
protected Expression visitWindowOperation(WindowOperation node, Context<C> context)
{
if (!context.isDefaultRewrite()) {
Expression result = rewriter.rewriteWindowOperation(node, context.get(), ExpressionTreeRewriter.this);
if (result != null) {
return result;
}
}

Identifier name = rewrite(node.getName(), context.get());
Window window = rewriteWindow(node.getWindow(), context);

if (name != node.getName() || window != node.getWindow()) {
return new WindowOperation(name, window);
}

return node;
}

@Override
protected Expression visitLambdaExpression(LambdaExpression node, Context<C> context)
{
Expand Down
Loading

0 comments on commit f98fe48

Please sign in to comment.