[CALCITE-129] Support recursive WITH queries

Close apache#3480

core/src/main/codegen/templates/Parser.jj

@@ -3486,14 +3486,16 @@ SqlNodeList WithList() :
 {
     final Span s;
     final List<SqlWithItem> list = new ArrayList<SqlWithItem>();
+    boolean recursive = false;
 }
 {
-    <WITH> { s = span(); }
-    AddWithItem(list) ( <COMMA> AddWithItem(list) )*
+    <WITH> [ <RECURSIVE> { recursive = true; } ]{ s = span(); }
+    AddWithItem(list, SqlLiteral.createBoolean(recursive, getPos()))
+    ( <COMMA> AddWithItem(list, SqlLiteral.createBoolean(recursive, getPos())) )*
     { return new SqlNodeList(list, s.end(this)); }
 }
 
-void AddWithItem(List<SqlWithItem> list) :
+void AddWithItem(List<SqlWithItem> list, SqlLiteral recursive) :
 {
     final SqlIdentifier id;
     final SqlNodeList columnList;
@@ -3504,7 +3506,7 @@ void AddWithItem(List<SqlWithItem> list) :
     ( columnList = ParenthesizedSimpleIdentifierList() | { columnList = null; } )
     <AS>
     definition = ParenthesizedExpression(ExprContext.ACCEPT_QUERY)
-    { list.add(new SqlWithItem(id.getParserPosition(), id, columnList, definition)); }
+    { list.add(new SqlWithItem(id.getParserPosition(), id, columnList, definition, recursive)); }
 }
 
 /**

core/src/main/java/org/apache/calcite/plan/RelOptUtil.java

@@ -3304,6 +3304,28 @@ public static RelNode createProject(RelNode child, Mappings.TargetMapping mappin
     return createProject(projectFactory, child, Mappings.asListNonNull(mapping.inverse()));
   }
 
+  /** Returns the relational table node for {@code tableName} if it occurs within a
+   * relational expression {@code root} otherwise an empty option is returned. */
+  public static @Nullable RelOptTable findTable(RelNode root, final String tableName) {
+    try {
+      RelShuttle visitor = new RelHomogeneousShuttle() {
+        @Override public RelNode visit(TableScan scan) {
+          final RelOptTable scanTable = scan.getTable();
+          final List<String> qualifiedName = scanTable.getQualifiedName();
+          if (qualifiedName.get(qualifiedName.size() - 1).equals(tableName)) {
+            throw new Util.FoundOne(scanTable);
+          }
+          return super.visit(scan);
+        }
+      };
+      root.accept(visitor);
+      return null;
+    } catch (Util.FoundOne e) {
+      Util.swallow(e, null);
+      return (RelOptTable) e.getNode();
+    }
+  }
+
   /** Returns whether relational expression {@code target} occurs within a
    * relational expression {@code ancestor}. */
   public static boolean contains(RelNode ancestor, final RelNode target) {

core/src/main/java/org/apache/calcite/runtime/CalciteResource.java

@@ -849,6 +849,12 @@ ExInst<CalciteException> illegalArgumentForTableFunctionCall(String a0,
   @BaseMessage("Must contain an ORDER BY clause when WITHIN is used")
   ExInst<SqlValidatorException> cannotUseWithinWithoutOrderBy();
 
+  @BaseMessage("A recursive query only supports UNION [ALL] operator")
+  ExInst<SqlValidatorException> recursiveWithMustHaveUnionSetOp();
+
+  @BaseMessage("A recursive query only supports binary UNION [ALL] operator")
+  ExInst<SqlValidatorException> recursiveWithMustHaveTwoChildUnionSetOp();
+
   @BaseMessage("First column of ORDER BY must be of type TIMESTAMP")
   ExInst<SqlValidatorException> firstColumnOfOrderByMustBeTimestamp();
 

core/src/main/java/org/apache/calcite/sql/SqlKind.java

@@ -195,6 +195,9 @@ public enum SqlKind {
   /** Item in WITH clause. */
   WITH_ITEM,
 
+  /** Represents a recursive CTE as a table ref. */
+  WITH_ITEM_TABLE_REF,
+
   /** Item expression. */
   ITEM,
 

core/src/main/java/org/apache/calcite/sql/SqlWith.java

@@ -96,6 +96,10 @@ private SqlWithOperator() {
       final SqlWith with = (SqlWith) call;
       final SqlWriter.Frame frame =
           writer.startList(SqlWriter.FrameTypeEnum.WITH, "WITH", "");
+      boolean isRecursive = ((SqlWithItem) with.withList.get(0)).recursive.booleanValue();
+      if (isRecursive) {
+        writer.keyword("RECURSIVE");
+      }
       final SqlWriter.Frame frame1 = writer.startList("", "");
       for (SqlNode node : with.withList) {
         writer.sep(",");

core/src/main/java/org/apache/calcite/sql/SqlWithItem.java

@@ -30,13 +30,23 @@
 public class SqlWithItem extends SqlCall {
   public SqlIdentifier name;
   public @Nullable SqlNodeList columnList; // may be null
+  public SqlLiteral recursive;
   public SqlNode query;
 
+  @Deprecated // to be removed before 2.0
   public SqlWithItem(SqlParserPos pos, SqlIdentifier name,
       @Nullable SqlNodeList columnList, SqlNode query) {
+    this(pos, name, columnList, query,
+        SqlLiteral.createBoolean(false, SqlParserPos.ZERO));
+  }
+
+  public SqlWithItem(SqlParserPos pos, SqlIdentifier name,
+      @Nullable SqlNodeList columnList, SqlNode query,
+      SqlLiteral recursive) {
     super(pos);
     this.name = name;
     this.columnList = columnList;
+    this.recursive = recursive;
     this.query = query;
   }
 
@@ -106,7 +116,8 @@ private static class SqlWithItemOperator extends SqlSpecialOperator {
       assert functionQualifier == null;
       assert operands.length == 3;
       return new SqlWithItem(pos, (SqlIdentifier) operands[0],
-          (SqlNodeList) operands[1], operands[2]);
+          (SqlNodeList) operands[1], operands[2],
+          SqlLiteral.createBoolean(false, SqlParserPos.ZERO));
     }
   }
 }

core/src/main/java/org/apache/calcite/sql/validate/SqlValidatorImpl.java

@@ -123,6 +123,7 @@
 
 import org.apiguardian.api.API;
 import org.checkerframework.checker.nullness.qual.KeyFor;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.checker.nullness.qual.PolyNull;
 import org.checkerframework.dataflow.qual.Pure;
@@ -2895,7 +2896,7 @@ private void registerQuery(
           parentScope,
           usingScope,
           node,
-          node,
+          enclosingNode,
           alias,
           forceNullable);
       break;
@@ -3099,9 +3100,23 @@ private void registerSetop(
 
     // A setop is in the same scope as its parent.
     scopes.put(call, parentScope);
-    for (SqlNode operand : call.getOperandList()) {
+    @NonNull SqlValidatorScope recursiveScope = parentScope;
+    if (enclosingNode.getKind() == SqlKind.WITH_ITEM) {
+      if (node.getKind() != SqlKind.UNION) {
+        throw newValidationError(node, RESOURCE.recursiveWithMustHaveUnionSetOp());
+      } else if (call.getOperandList().size() > 2) {
+        throw newValidationError(node, RESOURCE.recursiveWithMustHaveTwoChildUnionSetOp());
+      }
+      final WithScope scope = (WithScope) scopes.get(enclosingNode);
+      // recursive scope is only set for the recursive queries.
+      recursiveScope = scope != null && scope.recursiveScope != null
+          ? Objects.requireNonNull(scope.recursiveScope) : parentScope;
+    }
+    for (int i = 0; i < call.getOperandList().size(); i++) {
+      SqlNode operand = call.getOperandList().get(i);
+      @NonNull SqlValidatorScope scope = i == 0 ? parentScope : recursiveScope;
       registerQuery(
-          parentScope,
+          scope,
           null,
           operand,
           operand,
@@ -3126,17 +3141,21 @@ private void registerWith(
     SqlValidatorScope scope = parentScope;
     for (SqlNode withItem_ : with.withList) {
       final SqlWithItem withItem = (SqlWithItem) withItem_;
-      final WithScope withScope = new WithScope(scope, withItem);
+
+      final boolean isRecursiveWith = withItem.recursive.booleanValue();
+      final SqlValidatorScope withScope =
+          new WithScope(scope, withItem,
+              isRecursiveWith ? new WithRecursiveScope(scope, withItem) : null);
       scopes.put(withItem, withScope);
 
-      registerQuery(scope, null, withItem.query, with,
-          withItem.name.getSimple(), false);
+      registerQuery(scope, null, withItem.query,
+          withItem.recursive.booleanValue() ? withItem : with, withItem.name.getSimple(),
+          forceNullable);
       registerNamespace(null, alias,
           new WithItemNamespace(this, withItem, enclosingNode),
           false);
       scope = withScope;
     }
-
     registerQuery(scope, null, with.body, enclosingNode, alias, forceNullable,
         checkUpdate);
   }

core/src/main/java/org/apache/calcite/sql/validate/SqlWithItemTableRef.java

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.sql.validate;
+
+import org.apache.calcite.sql.SqlCall;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlLiteral;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlNodeList;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlSpecialOperator;
+import org.apache.calcite.sql.SqlTableRef;
+import org.apache.calcite.sql.SqlWithItem;
+import org.apache.calcite.sql.parser.SqlParserPos;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import static java.util.Objects.requireNonNull;
+
+/**
+ * A <code>SqlWithItemTableRef</code> is a node created during validation for
+ * recursive queries which represents a table reference in a {@code WITH RECURSIVE} clause.
+ */
+public class SqlWithItemTableRef extends SqlTableRef {
+  private final SqlWithItem withItem;
+  public SqlWithItemTableRef(SqlParserPos pos,
+      SqlWithItem withItem) {
+    super(pos, withItem.name, SqlNodeList.EMPTY);
+    this.withItem = withItem;
+  }
+
+  private static final SqlOperator OPERATOR =
+      new SqlSpecialOperator("WITH_ITEM_TABLE_REF", SqlKind.WITH_ITEM_TABLE_REF) {
+        @Override public SqlCall createCall(
+            @Nullable SqlLiteral functionQualifier,
+            SqlParserPos pos, @Nullable SqlNode... operands) {
+          return new SqlWithItemTableRef(pos,
+              (SqlWithItem) requireNonNull(operands[0], "withItem"));
+        }
+      };
+  @Override public SqlOperator getOperator() {
+    return OPERATOR;
+  }
+
+  public SqlWithItem getWithItem() {
+    return withItem;
+  }
+}

core/src/main/java/org/apache/calcite/sql/validate/WithItemNamespace.java

@@ -32,14 +32,14 @@ class WithItemNamespace extends AbstractNamespace {
   private final SqlWithItem withItem;
 
   WithItemNamespace(SqlValidatorImpl validator, SqlWithItem withItem,
-      SqlNode enclosingNode) {
+      @Nullable SqlNode enclosingNode) {
     super(validator, enclosingNode);
     this.withItem = withItem;
   }
 
   @Override protected RelDataType validateImpl(RelDataType targetRowType) {
     final SqlValidatorNamespace childNs =
-        validator.getNamespaceOrThrow(withItem.query);
+        validator.getNamespaceOrThrow(getQuery());
     final RelDataType rowType = childNs.getRowTypeSansSystemColumns();
     SqlNodeList columnList = withItem.columnList;
     if (columnList == null) {
@@ -53,6 +53,12 @@ class WithItemNamespace extends AbstractNamespace {
     return builder.build();
   }
 
+  /** Returns the node from which {@link #validateImpl(RelDataType)} determines
+   * the namespace. */
+  protected SqlNode getQuery() {
+    return withItem.query;
+  }
+
   @Override public @Nullable SqlNode getNode() {
     return withItem;
   }

core/src/main/java/org/apache/calcite/sql/validate/WithItemRecursiveNamespace.java

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.sql.validate;
+
+import org.apache.calcite.sql.SqlCall;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlWith;
+import org.apache.calcite.sql.SqlWithItem;
+import org.apache.calcite.sql.parser.SqlParserPos;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/** Very similar to {@link WithItemNamespace} but created only for RECURSIVE queries. */
+class WithItemRecursiveNamespace extends WithItemNamespace {
+  private final SqlWithItem withItem;
+  private final SqlWithItemTableRef withItemTableRef;
+
+  /**
+   * Creates a Namespace for a query specified in {@code WITH RECURSIVE} clause.
+   *
+   * @param validator Validator
+   * @param withItem A with query item specified in {@code WITH} clause
+   * @param enclosingNode Enclosing node
+   */
+  WithItemRecursiveNamespace(SqlValidatorImpl validator,
+      SqlWithItem withItem,
+      @Nullable SqlNode enclosingNode) {
+    super(validator, withItem, enclosingNode);
+    this.withItem = withItem;
+    this.withItemTableRef = new SqlWithItemTableRef(SqlParserPos.ZERO, withItem);
+  }
+
+  @Override protected SqlNode getQuery() {
+    SqlNode call = this.withItem.query;
+    while (call.getKind() == SqlKind.WITH) {
+      call = ((SqlWith) call).body;
+    }
+    return ((SqlCall) call).operand(0);
+  }
+
+  @Override public @Nullable SqlNode getNode() {
+    return withItemTableRef;
+  }
+}