Preserve LATERAL UNNEST during SqlNode validation (#101)

* preserve lateral unnest

* preserve lateral unnest unit test

* modify conditional

* only preserve lateral, do not addd

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

@@ -31,6 +31,7 @@
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.type.SqlTypeUtil;
 import org.apache.calcite.sql.util.SqlBasicVisitor;
+import org.apache.calcite.sql.util.SqlShuttle;
 import org.apache.calcite.sql.validate.SqlNameMatcher;
 import org.apache.calcite.sql.validate.SqlValidatorUtil;
 import org.apache.calcite.util.BarfingInvocationHandler;
@@ -1079,6 +1080,17 @@ private void visitChild(SqlNode node) {
       return check(type);
     }
   }
+
+  /** Walks over a {@link org.apache.calcite.sql.SqlNode} tree and removes any
+   * instance of a LATERAL operator. */
+  public static class LateralRemover extends SqlShuttle {
+    @Override public SqlNode visit(SqlCall call) {
+      if (call.getKind() == SqlKind.LATERAL) {
+        return call.operand(0).accept(this);
+      }
+      return super.visit(call);
+    }
+  }
 }
 
 // End SqlUtil.java

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

@@ -2006,6 +2006,23 @@ protected void registerNamespace(
     }
   }
 
+  /**
+   * @see #registerFrom(SqlValidatorScope, SqlValidatorScope, boolean, SqlNode, SqlNode, String, SqlNodeList, boolean, boolean, boolean)
+   */
+  private SqlNode registerFrom(
+      SqlValidatorScope parentScope,
+      SqlValidatorScope usingScope,
+      boolean register,
+      final SqlNode node,
+      SqlNode enclosingNode,
+      String alias,
+      SqlNodeList extendList,
+      boolean forceNullable,
+      final boolean lateral) {
+    return registerFrom(parentScope, usingScope, register, node, enclosingNode, alias,
+        extendList, forceNullable, lateral, false);
+  }
+
   /**
    * Registers scopes and namespaces implied a relational expression in the
    * FROM clause.
@@ -2035,6 +2052,9 @@ protected void registerNamespace(
    * @param lateral       Whether LATERAL is specified, so that items to the
    *                      left of this in the JOIN tree are visible in the
    *                      scope
+   * @param fromLateral   Whether node was nested in a LATERAL node because
+   *                      even if LATERAL is specified, there may not be an explicit
+   *                      LATERAL node in the tree (e.g. UNNEST is implicitly LATERAL)
    * @return registered node, usually the same as {@code node}
    */
   private SqlNode registerFrom(
@@ -2046,7 +2066,8 @@ private SqlNode registerFrom(
       String alias,
       SqlNodeList extendList,
       boolean forceNullable,
-      final boolean lateral) {
+      final boolean lateral,
+      final boolean fromLateral) {
     final SqlKind kind = node.getKind();
 
     SqlNode expr;
@@ -2248,6 +2269,7 @@ private SqlNode registerFrom(
           alias,
           extendList,
           forceNullable,
+          true,
           true);
 
     case COLLECTION_TABLE:
@@ -2293,8 +2315,12 @@ private SqlNode registerFrom(
           alias,
           forceNullable);
 
-      // Keep preceding "LATERAL" keyword from joins with inner SELECT subquery
-      if (lateral && kind == SqlKind.SELECT) {
+      // Preserve "LATERAL" keyword in validated node:
+      // 1. Joins with inner SELECT subquery preceded by LATERAL
+      //    Example: SELECT ... FROM LATERAL(SELECT ...) AS t
+      // 2. Joins with UNNEST operator preceded by LATERAL
+      //    Example: SELECT ... FROM LATERAL(UNNEST ...) AS t
+      if (fromLateral && (kind == SqlKind.SELECT || kind == SqlKind.UNNEST)) {
         SqlNode lateralNode =
             SqlStdOperatorTable.LATERAL.createCall(POS, newNode);
         return lateralNode;
@@ -3120,6 +3146,7 @@ protected void validateFrom(
       SqlValidatorScope scope) {
     Objects.requireNonNull(targetRowType);
     switch (node.getKind()) {
+    case LATERAL:
     case AS:
       validateFrom(
           ((SqlCall) node).operand(0),
@@ -3138,24 +3165,14 @@ protected void validateFrom(
     case UNNEST:
       validateUnnest((SqlCall) node, scope, targetRowType);
       break;
-    case LATERAL:
-      // Validate subquery that LATERAL precedes
-      validateQuery(((SqlCall) node).operand(0), scope, targetRowType);
-      break;
     default:
       validateQuery(node, scope, targetRowType);
       break;
     }
 
     // Validate the namespace representation of the node, just in case the
     // validation did not occur implicitly.
-    if (node.getKind() == SqlKind.LATERAL) {
-      // Skip over fetching for LATERAL namespace since they aren't registered, use subquery instead
-      getNamespace(((SqlCall) node).operand(0), scope)
-          .validate(targetRowType);
-    } else {
-      getNamespace(node, scope).validate(targetRowType);
-    }
+    getNamespace(node, scope).validate(targetRowType);
   }
 
   protected void validateOver(SqlCall call, SqlValidatorScope scope) {

core/src/main/java/org/apache/calcite/sql2rel/SqlToRelConverter.java

@@ -1993,18 +1993,15 @@ protected void convertFrom(
       return;
     }
 
-    // Skip over LATERAL conversion
-    if (from.getKind() == SqlKind.LATERAL) {
-      from = ((SqlCall) from).operand(0);
-    }
-
     final SqlCall call;
     final SqlNode[] operands;
     switch (from.getKind()) {
     case MATCH_RECOGNIZE:
       convertMatchRecognize(bb, (SqlCall) from);
       return;
-
+    case LATERAL:
+      convertFrom(bb, ((SqlCall) from).operand(0));
+      return;
     case AS:
       call = (SqlCall) from;
       convertFrom(bb, call.operand(0));
@@ -2080,8 +2077,12 @@ protected void convertFrom(
               ((DelegatingScope) bb.scope).getParent());
       final Blackboard leftBlackboard =
           createBlackboard(leftScope, null, false);
+      // We don't register the scopes of LATERAL SqlNodes, so we need to
+      // remove them before looking up the associated scope
+      // Note that LATERALs could only appear on the right side of a join
+      SqlNode rightWithoutLaterals = right.accept(new SqlUtil.LateralRemover());
       final SqlValidatorScope rightScope =
-          Util.first(validator.getJoinScope(right),
+          Util.first(validator.getJoinScope(rightWithoutLaterals),
               ((DelegatingScope) bb.scope).getParent());
       final Blackboard rightBlackboard =
           createBlackboard(rightScope, null, false);

core/src/test/java/org/apache/calcite/test/SqlValidatorTest.java

@@ -11521,7 +11521,7 @@ private void checkCustomColumnResolving(String table) {
    * <a href="https://github.com/linkedin/linkedin-calcite/pull/98">
    * Preserve LATERAL keyword during validation #98</a>. */
   @Test
-  public void testLateralKeywordExistsAfterValidation() throws SqlParseException {
+  public void testLateralPreservedInLateralSelect() throws SqlParseException {
     String sql = "SELECT * FROM emp CROSS JOIN LATERAL "
         + "(SELECT * FROM dept WHERE deptno = emp.deptno)";
 
@@ -11531,6 +11531,24 @@ public void testLateralKeywordExistsAfterValidation() throws SqlParseException {
 
     assertTrue(validatedNode.toString().contains("LATERAL"));
   }
+
+  @Test
+  public void testLateralKeywordIsNotAddedToUnnest() throws SqlParseException {
+    // Per SQL std, UNNEST is implicitly LATERAL
+    String sql = "select*from unnest(array[1])";
+    SqlNode node = tester.parseQuery(sql);
+    SqlValidator validator = tester.getValidator();
+    SqlNode validatedNode = validator.validate(node);
+
+    assertTrue(!validatedNode.toString().contains("LATERAL"));
+
+    sql = "select c from unnest(array(select deptno from dept)) as t(c)";
+    node = tester.parseQuery(sql);
+    validator = tester.getValidator();
+    validatedNode = validator.validate(node);
+
+    assertTrue(!validatedNode.toString().contains("LATERAL"));
+  }
 }
 
 // End SqlValidatorTest.java