Multi-Stage: Add SQL:2016 MATCH_RECOGNIZE (row pattern recognition)

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -663,6 +663,18 @@ public static Integer getMaxRowsInWindow(Map<String, String> queryOptions) {
     return checkedParseIntPositive(QueryOptionKey.MAX_ROWS_IN_WINDOW, maxRowsInWindow);
   }
 
+  @Nullable
+  public static Integer getMaxRowsInMatchPartition(Map<String, String> queryOptions) {
+    String maxRowsInMatchPartition = queryOptions.get(QueryOptionKey.MAX_ROWS_IN_MATCH_PARTITION);
+    return checkedParseIntPositive(QueryOptionKey.MAX_ROWS_IN_MATCH_PARTITION, maxRowsInMatchPartition);
+  }
+
+  @Nullable
+  public static Long getMaxStepsPerMatchAttempt(Map<String, String> queryOptions) {
+    String maxStepsPerMatchAttempt = queryOptions.get(QueryOptionKey.MAX_STEPS_PER_MATCH_ATTEMPT);
+    return checkedParseLongPositive(QueryOptionKey.MAX_STEPS_PER_MATCH_ATTEMPT, maxStepsPerMatchAttempt);
+  }
+
   @Nullable
   public static WindowOverFlowMode getWindowOverflowMode(Map<String, String> queryOptions) {
     String windowOverflowModeStr = queryOptions.get(QueryOptionKey.WINDOW_OVERFLOW_MODE);

pinot-common/src/main/proto/expressions.proto

@@ -52,6 +52,30 @@ message InputRef {
   int32 index = 1;
 }
 
+// A reference to a column of the row bound to a specific MATCH_RECOGNIZE pattern variable, as it appears inside
+// MEASURES and DEFINE expressions. Mirrors Calcite's RexPatternFieldRef, which carries the pattern variable
+// alongside the input index.
+//
+// This is deliberately NOT modelled as an InputRef. RexPatternFieldRef extends RexInputRef, so anything that only
+// looks at the index silently turns `DEFINE UP AS UP.price > PREV(UP.price)` into a read of the current row's
+// column: results that are wrong but still type-correct, and therefore invisible. Any consumer that does not
+// understand this expression case must fail loudly instead of degrading it to an InputRef.
+//
+// Only meaningful inside a MatchNode: symbolOrdinal indexes MatchNode.patternSymbols.
+message PatternFieldRef {
+  // Column index into the MatchNode input row.
+  int32 index = 1;
+  // Ordinal of the pattern variable this reference is bound to, i.e. the index into MatchNode.patternSymbols.
+  // The sentinel -2 denotes the SQL:2016 universal row pattern variable: an unqualified reference that ranges over
+  // every row in the match regardless of its pattern variable. -1 is reserved for an unresolved planner-local
+  // reference and must never reach the wire; values below -2 are invalid. This field is authoritative.
+  int32 symbolOrdinal = 2;
+  // Pattern variable name as written in the query, carried for explain plans and error messages only. For the
+  // universal ordinal (-2), Calcite supplies the row-source/table alias here; it is not a pattern-variable name.
+  // Consumers must resolve the variable through symbolOrdinal and never string-match on this field.
+  string alpha = 3;
+}
+
 message Literal {
   ColumnDataType dataType = 1;
   oneof value {
@@ -108,5 +132,6 @@ message Expression {
     InputRef inputRef = 1;
     Literal literal = 2;
     FunctionCall functionCall = 3;
+    PatternFieldRef patternFieldRef = 4;
   }
 }

pinot-common/src/main/proto/plan.proto

@@ -44,6 +44,10 @@ message PlanNode {
     // serialized by an older-version broker can still be deserialized. No current broker sets this field.
     EnrichedJoinNode enrichedJoinNode = 17 [deprecated = true];
     UnnestNode unnestNode = 18;
+    // Rolling-upgrade boundary: servers released before MATCH_RECOGNIZE do not know field 19. Protobuf preserves it
+    // as an unknown field, but their PlanNodeDeserializer sees NODE_NOT_SET and cannot execute the stage. Upgrade all
+    // servers before issuing MATCH_RECOGNIZE queries; there is no mixed-version fallback to another node kind.
+    MatchNode matchNode = 19;
   }
 }
 
@@ -264,6 +268,119 @@ message UnnestNode {
   bool prunedPassthrough = 6;
 }
 
+// ---------------------------------------------------------------------------------------------------------------
+// MATCH_RECOGNIZE (SQL:2016 row pattern recognition)
+// ---------------------------------------------------------------------------------------------------------------
+
+// SQL:2016 `AFTER MATCH SKIP` clause: where pattern matching resumes after a successful match.
+//
+// SKIP_PAST_LAST_ROW is 0 on purpose. It is the SQL:2016 default (and the default in Trino, Snowflake and Oracle),
+// so a MatchNode that omits the field falls back to non-overlapping matches rather than to the overlapping
+// SKIP TO NEXT ROW semantics that Calcite defaults to internally when the clause is omitted.
+//
+// The values are wire-stable; renumbering them silently changes query results across mixed-version brokers/servers.
+enum AfterMatchSkipMode {
+  SKIP_PAST_LAST_ROW = 0;
+  SKIP_TO_NEXT_ROW = 1;
+  SKIP_TO_FIRST = 2;
+  SKIP_TO_LAST = 3;
+}
+
+// SQL:2016 `ONE ROW PER MATCH` / `ALL ROWS PER MATCH`. ALL_ROWS_PER_MATCH is declared so that its wire value is
+// pinned, but it is rejected at planning time until per-row output is implemented.
+enum RowsPerMatchMode {
+  ONE_ROW_PER_MATCH = 0;
+  ALL_ROWS_PER_MATCH = 1;
+}
+
+// The kind of a RowPattern tree node. PATTERN_EXCLUDE and PATTERN_PERMUTE are declared up front so the deferred
+// work cannot renumber the supported kinds; a plan carrying them is rejected at deserialization time.
+enum RowPatternKind {
+  // A single pattern variable. Carries symbolOrdinal.
+  PATTERN_SYMBOL = 0;
+  // Juxtaposition, e.g. `A B C`. Two or more children, in pattern order.
+  PATTERN_CONCAT = 1;
+  // Alternation `|`, e.g. `A | B`. Two or more children, in pattern order (leftmost alternative wins).
+  PATTERN_ALTERNATE = 2;
+  // A quantifier applied to exactly one child. Carries quantifier.
+  PATTERN_QUANTIFIER = 3;
+  // The `^` anchor: the match must start at the first row of the partition. No children.
+  PATTERN_ANCHOR_START = 4;
+  // The `$` anchor: the match must end at the last row of the partition. No children.
+  PATTERN_ANCHOR_END = 5;
+  // Deferred: `{- ... -}` exclusion. Exactly one child.
+  PATTERN_EXCLUDE = 6;
+  // Deferred: `PERMUTE(...)`. Two or more children.
+  PATTERN_PERMUTE = 7;
+}
+
+// A `{n,m}` style repetition applied to the single child of a PATTERN_QUANTIFIER node. `*`, `+` and `?` are
+// normalized into this form: `*` is {0,-1}, `+` is {1,-1}, `?` is {0,1}.
+message RowPatternQuantifier {
+  // Minimum number of repetitions, >= 0.
+  int32 minRepeat = 1;
+  // Maximum number of repetitions, or -1 for unbounded (`*`, `+`, `{n,}`).
+  int32 maxRepeat = 2;
+  // False for the reluctant forms (`*?`, `+?`, `??`, `{n,m}?`), which prefer the shortest match.
+  bool greedy = 3;
+}
+
+// One node of the row pattern tree. This is a self-contained representation of the PATTERN clause: pattern
+// variables are referenced by their ordinal in MatchNode.patternSymbols, so a consumer never string-matches
+// variable names.
+message RowPattern {
+  RowPatternKind kind = 1;
+  // PATTERN_SYMBOL only: ordinal into MatchNode.patternSymbols.
+  int32 symbolOrdinal = 2;
+  // PATTERN_CONCAT / PATTERN_ALTERNATE / PATTERN_PERMUTE: two or more children, in pattern order.
+  // PATTERN_QUANTIFIER / PATTERN_EXCLUDE: exactly one child. Empty for PATTERN_SYMBOL and the anchors.
+  repeated RowPattern children = 3;
+  // PATTERN_QUANTIFIER only.
+  RowPatternQuantifier quantifier = 4;
+}
+
+// One entry of the pattern variable symbol table. The variable's ordinal is its index in
+// MatchNode.patternSymbols.
+message PatternSymbol {
+  // Pattern variable name as written in PATTERN / DEFINE. Informational: consumers resolve variables by ordinal.
+  string name = 1;
+  // The DEFINE predicate for this variable. Absent when the variable has no DEFINE entry, which per SQL:2016
+  // means it matches every row.
+  optional Expression definition = 2;
+}
+
+// One `MEASURES <expression> AS <name>` item.
+message MatchMeasure {
+  // Output column name.
+  string name = 1;
+  Expression expression = 2;
+}
+
+message MatchNode {
+  // Pattern variable symbol table. A symbol's ordinal is its index in this list; every PATTERN_SYMBOL node and
+  // every PatternFieldRef refers to a variable by that ordinal.
+  repeated PatternSymbol patternSymbols = 1;
+  // Root of the row pattern tree (the PATTERN clause).
+  RowPattern pattern = 2;
+  // MEASURES items, in output order.
+  repeated MatchMeasure measures = 3;
+  // PARTITION BY: column indexes into the input row.
+  repeated int32 partitionKeys = 4;
+  // ORDER BY. Mandatory for MATCH_RECOGNIZE, so this is never empty.
+  repeated Collation collations = 5;
+  AfterMatchSkipMode afterMatchSkipMode = 6;
+  // For SKIP_TO_FIRST / SKIP_TO_LAST: ordinal of the target pattern variable. Absent for the other skip modes.
+  // Explicitly optional because ordinal 0 is a valid variable, so the proto3 default cannot double as "unset".
+  optional int32 afterMatchSkipToSymbolOrdinal = 7;
+  RowsPerMatchMode rowsPerMatchMode = 8;
+
+  // Reserved for the deferred SQL:2016 clauses so that implementing them later cannot renumber anything above:
+  //    9 -> SUBSET: `repeated PatternSubset subsets` (union variable name -> member symbol ordinals)
+  //   10 -> WITHIN: the match time interval
+  reserved 9, 10;
+  reserved "subsets", "within";
+}
+
 enum WindowFrameType {
   ROWS = 0;
   RANGE = 1;

pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java

@@ -38,7 +38,7 @@
 public class QueryOptionsUtilsTest {
   private static final List<String> POSITIVE_INT_KEYS =
       List.of(NUM_REPLICA_GROUPS_TO_QUERY, MAX_EXECUTION_THREADS, NUM_GROUPS_LIMIT, MAX_INITIAL_RESULT_HOLDER_CAPACITY,
-          MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW);
+          MAX_STREAMING_PENDING_BLOCKS, MAX_ROWS_IN_JOIN, MAX_ROWS_IN_WINDOW, MAX_ROWS_IN_MATCH_PARTITION);
   private static final List<String> NON_NEGATIVE_INT_KEYS = List.of(MULTI_STAGE_LEAF_LIMIT);
   private static final List<String> UNBOUNDED_INT_KEYS =
       List.of(MIN_SEGMENT_GROUP_TRIM_SIZE, MIN_SERVER_GROUP_TRIM_SIZE, MIN_BROKER_GROUP_TRIM_SIZE,
@@ -49,7 +49,8 @@ public class QueryOptionsUtilsTest {
       addAll(UNBOUNDED_INT_KEYS);
     }};
   private static final List<String> POSITIVE_LONG_KEYS =
-      List.of(TIMEOUT_MS, MAX_SERVER_RESPONSE_SIZE_BYTES, MAX_QUERY_RESPONSE_SIZE_BYTES);
+      List.of(TIMEOUT_MS, MAX_SERVER_RESPONSE_SIZE_BYTES, MAX_QUERY_RESPONSE_SIZE_BYTES,
+          MAX_STEPS_PER_MATCH_ATTEMPT);
 
   @Test
   public void shouldConvertCaseInsensitiveMapToUseCorrectValues() {
@@ -333,6 +334,8 @@ private static Object getValue(Map<String, String> map, String key) {
         return QueryOptionsUtils.getMaxRowsInJoin(map);
       case MAX_ROWS_IN_WINDOW:
         return QueryOptionsUtils.getMaxRowsInWindow(map);
+      case MAX_ROWS_IN_MATCH_PARTITION:
+        return QueryOptionsUtils.getMaxRowsInMatchPartition(map);
       // Non-negative ints
       case MULTI_STAGE_LEAF_LIMIT:
         return QueryOptionsUtils.getMultiStageLeafLimit(map);
@@ -352,6 +355,8 @@ private static Object getValue(Map<String, String> map, String key) {
         return QueryOptionsUtils.getMaxServerResponseSizeBytes(map);
       case MAX_QUERY_RESPONSE_SIZE_BYTES:
         return QueryOptionsUtils.getMaxQueryResponseSizeBytes(map);
+      case MAX_STEPS_PER_MATCH_ATTEMPT:
+        return QueryOptionsUtils.getMaxStepsPerMatchAttempt(map);
       default:
         throw new IllegalArgumentException("Unexpected key!");
     }

pinot-common/src/test/java/org/apache/pinot/sql/parsers/SqlQueryOptionValidationTest.java

@@ -30,6 +30,7 @@
 import org.apache.logging.log4j.core.config.Property;
 import org.apache.pinot.common.utils.config.QueryOptionsUtils;
 import org.apache.pinot.common.utils.config.QueryOptionsUtils.SqlQueryOptionValidationMode;
+import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.Test;
 
@@ -108,6 +109,18 @@ public void rejectModeAcceptsKnownKeysCaseInsensitivelyAndTraceAndDatabase() {
     assertEquals(options.get("Database"), "db1");
   }
 
+  @Test
+  public void matchRecognizeOptionsAreKnownAndCanonicalizedCaseInsensitively() {
+    QueryOptionsUtils.setSqlQueryOptionValidationMode(SqlQueryOptionValidationMode.REJECT);
+
+    Map<String, String> options = optionsOf("SET MAXROWSINMATCHPARTITION='100'; "
+        + "SET maxstepspermatchattempt='200'; SET ALLOWMATCHRECOGNIZEWITHOUTPARTITIONBY='true'; "
+        + "select * from vegetables");
+    assertEquals(options.get(QueryOptionKey.MAX_ROWS_IN_MATCH_PARTITION), "100");
+    assertEquals(options.get(QueryOptionKey.MAX_STEPS_PER_MATCH_ATTEMPT), "200");
+    assertEquals(options.get(QueryOptionKey.ALLOW_MATCH_RECOGNIZE_WITHOUT_PARTITION_BY), "true");
+  }
+
   @Test
   public void rejectModeLeavesDmlOptionsFreeForm() {
     QueryOptionsUtils.setSqlQueryOptionValidationMode(SqlQueryOptionValidationMode.REJECT);
@@ -165,7 +178,9 @@ public void warnModePreservesUnknownKeysAndLogsOncePerDistinctKey() {
 
       // Known keys are never logged.
       int loggedSoFar = appender.messagesContaining("Unsupported query option").size();
-      optionsOf("SET timeoutMs='100'; select * from vegetables");
+      optionsOf("SET timeoutMs='100'; SET maxRowsInMatchPartition='200'; "
+          + "SET maxStepsPerMatchAttempt='300'; SET allowMatchRecognizeWithoutPartitionBy='true'; "
+          + "select * from vegetables");
       assertEquals(appender.messagesContaining("Unsupported query option").size(), loggedSoFar);
     } finally {
       appender.detach();