Skip to content

Commit

Permalink
Support subqueries in pattern matching context in window
Browse files Browse the repository at this point in the history
  • Loading branch information
@kasiafi
kasiafi committed Aug 6, 2021
1 parent 3318d46 commit 236a6d4
Show file tree
Hide file tree
Showing 2 changed files with 105 additions and 1 deletion.
23 changes: 22 additions & 1 deletion core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@
import io.trino.sql.tree.LambdaArgumentDeclaration;
import io.trino.sql.tree.LambdaExpression;
import io.trino.sql.tree.LongLiteral;
import io.trino.sql.tree.MeasureDefinition;
import io.trino.sql.tree.Node;
import io.trino.sql.tree.NodeRef;
import io.trino.sql.tree.Offset;
Expand All @@ -87,6 +88,7 @@
import io.trino.sql.tree.Table;
import io.trino.sql.tree.Union;
import io.trino.sql.tree.Update;
import io.trino.sql.tree.VariableDefinition;
import io.trino.sql.tree.WindowFrame;
import io.trino.sql.tree.WindowOperation;
import io.trino.type.TypeCoercion;
Expand Down Expand Up @@ -1016,6 +1018,8 @@ else if (window.getFrame().isPresent()) {
}

if (window.getFrame().isPresent() && window.getFrame().get().getPattern().isPresent()) {
WindowFrame frame = window.getFrame().get();
subPlan = subqueryPlanner.handleSubqueries(subPlan, extractPatternRecognitionExpressions(frame.getVariableDefinitions(), frame.getMeasures()), analysis.getSubqueries(node));
subPlan = planPatternRecognition(subPlan, windowFunction, window, coercions, frameEnd);
}
else {
Expand Down Expand Up @@ -1413,7 +1417,8 @@ private PlanBuilder planWindowMeasures(Node node, PlanBuilder subPlan, List<Wind
.addAll(getSortItemsFromOrderBy(window.getOrderBy()).stream()
.map(SortItem::getSortKey)
.iterator());
Optional<Expression> endValue = window.getFrame().get().getEnd().get().getValue();
WindowFrame frame = window.getFrame().get();
Optional<Expression> endValue = frame.getEnd().get().getValue();
endValue.ifPresent(inputsBuilder::add);

List<Expression> inputs = inputsBuilder.build();
Expand All @@ -1426,12 +1431,28 @@ private PlanBuilder planWindowMeasures(Node node, PlanBuilder subPlan, List<Wind
subPlan = plan.getSubPlan();
Optional<Symbol> frameEnd = plan.getFrameOffsetSymbol();

subPlan = subqueryPlanner.handleSubqueries(subPlan, extractPatternRecognitionExpressions(frame.getVariableDefinitions(), frame.getMeasures()), analysis.getSubqueries(node));
subPlan = planPatternRecognition(subPlan, windowMeasure, window, frameEnd);
}

return subPlan;
}

private static List<Expression> extractPatternRecognitionExpressions(List<VariableDefinition> variableDefinitions, List<MeasureDefinition> measureDefinitions)
{
ImmutableList.Builder<Expression> expressions = ImmutableList.builder();

variableDefinitions.stream()
.map(VariableDefinition::getExpression)
.forEach(expressions::add);

measureDefinitions.stream()
.map(MeasureDefinition::getExpression)
.forEach(expressions::add);

return expressions.build();
}

private PlanBuilder planPatternRecognition(
PlanBuilder subPlan,
WindowOperation windowMeasure,
Expand Down
83 changes: 83 additions & 0 deletions core/trino-main/src/test/java/io/trino/sql/query/TestRowPatternMatchingInWindow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1396,4 +1396,87 @@ public void testRowPatternMatchingInOrderBy()
" ) ASC NULLS FIRST"))
.matches("VALUES 6, 3, 4, 5, 1, 2");
}

@Test
public void testSubqueries()
{
// NOTE: Subquery in pattern recognition context can always be handled prior to pattern recognition planning, because
// by analysis it cannot contain labeled column references, navigations, `CLASSIFIER()`, or `MATCH_NUMBER()` calls,
// which mustn't be pre-projected.
String query = "SELECT id, val OVER w " +
" FROM (VALUES " +
" (1, 100), " +
" (2, 200), " +
" (3, 300), " +
" (4, 400) " +
" ) t(id, value) " +
" WINDOW w AS ( " +
" ORDER BY id " +
" MEASURES %s AS val " +
" ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING " +
" AFTER MATCH SKIP TO NEXT ROW " +
" PATTERN (A+) " +
" DEFINE A AS %s " +
" )";

assertThat(assertions.query(format(query, "(SELECT 'x')", "(SELECT true)")))
.matches("VALUES " +
" (1, 'x'), " +
" (2, 'x'), " +
" (3, 'x'), " +
" (4, 'x') ");

// subquery nested in navigation
assertThat(assertions.query(format(query, "LAST(A.value + (SELECT 1000))", "FIRST(A.value < 0 OR (SELECT true))")))
.matches("VALUES " +
" (1, 1400), " +
" (2, 1400), " +
" (3, 1400), " +
" (4, 1400) ");

// IN-predicate: value and value list without column references
assertThat(assertions.query(format(query, "LAST(A.id < 0 OR 1 IN (SELECT 1))", "FIRST(A.id > 0 AND 1 IN (SELECT 1))")))
.matches("VALUES " +
" (1, true), " +
" (2, true), " +
" (3, true), " +
" (4, true) ");

// IN-predicate: unlabeled column reference in value
assertThat(assertions.query(format(query, "LAST(id + 50 IN (SELECT 100))", "LAST(id NOT IN (SELECT 100))")))
.matches("VALUES " +
" (1, false), " +
" (2, false), " +
" (3, false), " +
" (4, false) ");

// EXISTS-predicate
assertThat(assertions.query(format(query, "LAST(A.value < 0 OR EXISTS(SELECT 1))", "FIRST(A.value < 0 OR EXISTS(SELECT 1))")))
.matches("VALUES " +
" (1, true), " +
" (2, true), " +
" (3, true), " +
" (4, true) ");

// no pattern measures
assertThat(assertions.query("SELECT id, max(value) OVER w " +
" FROM (VALUES " +
" (1, 100), " +
" (2, 200), " +
" (3, 300), " +
" (4, 400) " +
" ) t(id, value) " +
" WINDOW w AS ( " +
" ORDER BY id " +
" ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING " +
" AFTER MATCH SKIP TO NEXT ROW " +
" PATTERN (A+) " +
" DEFINE A AS (SELECT true) " +
" )"))
.matches("VALUES " +
" (1, 400), " +
" (2, 400), " +
" (3, 400), " +
" (4, 400) ");
}
}

0 comments on commit 236a6d4

Please sign in to comment.