PullupCorrelatedPredicates should not produce unresolved plans.

Author: viirya

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/predicates.scala

@@ -140,25 +140,58 @@ case class In(value: Expression, list: Seq[Expression]) extends Predicate {
   require(list != null, "list should not be null")
   override def checkInputDataTypes(): TypeCheckResult = {
     list match {
-      case ListQuery(sub, _, _) :: Nil =>
+      case (l @ ListQuery(sub, children, _)) :: Nil =>
         val valExprs = value match {
           case cns: CreateNamedStruct => cns.valExprs
           case expr => Seq(expr)
         }
-        if (valExprs.length != sub.output.length) {
-          TypeCheckResult.TypeCheckFailure(
+
+        // SPARK-21759:
+        // It is possibly that the subquery plan has more output than value expressions, because
+        // the condition expressions in `ListQuery` might use part of subquery plan's output.
+        // For example, in the following query plan, the condition of `ListQuery` uses value#207
+        // from the subquery query. For now the size of output of subquery is 2, the size of value
+        // is 1.
+        //
+        //   Filter key#201 IN (list#200 [(value#207 = min(value)#204)])
+        //   :  +- Project [key#206, value#207]
+        //   :     +- Filter (value#207 > val_9)
+
+        // Take the subset of output which are not going to match with value expressions and also
+        // not used in condition expressions, if any.
+        val subqueryOutputNotInCondition = sub.output.drop(valExprs.length).filter { attr =>
+          l.children.forall { c =>
+            !c.references.contains(attr)
+          }
+        }
+
+        val basicErrorMessage =
+          s"""
+             |The number of columns in the left hand side of an IN subquery does not match the
+             |number of columns in the output of subquery.
+             |#columns in left hand side: ${valExprs.length}.
+             |#columns in right hand side: ${sub.output.length}.
+             |Left side columns:
+             |[${valExprs.map(_.sql).mkString(", ")}].
+             |Right side columns:
+             |[${sub.output.map(_.sql).mkString(", ")}].
+           """.stripMargin
+
+        if (valExprs.length > sub.output.length) {
+          TypeCheckResult.TypeCheckFailure(basicErrorMessage)
+        } else if (subqueryOutputNotInCondition.nonEmpty) {
+          val finalErrorMessage = basicErrorMessage +
             s"""
-               |The number of columns in the left hand side of an IN subquery does not match the
-               |number of columns in the output of subquery.
-               |#columns in left hand side: ${valExprs.length}.
-               |#columns in right hand side: ${sub.output.length}.
-               |Left side columns:
-               |[${valExprs.map(_.sql).mkString(", ")}].
-               |Right side columns:
-               |[${sub.output.map(_.sql).mkString(", ")}].
-             """.stripMargin)
+               | The additional output in subquery aren't used in the condition of subquery.
+               | Additional output:
+               | [${subqueryOutputNotInCondition.map(_.sql).mkString(", ")}].
+               | Condition:
+               | [${children.map(_.sql).mkString(", ")}].
+               | ${l.references}
+             """.stripMargin
+          TypeCheckResult.TypeCheckFailure(finalErrorMessage)
         } else {
-          val mismatchedColumns = valExprs.zip(sub.output).flatMap {
+          val mismatchedColumns = valExprs.zip(sub.output.take(valExprs.length)).flatMap {
             case (l, r) if l.dataType != r.dataType =>
               s"(${l.sql}:${l.dataType.catalogString}, ${r.sql}:${r.dataType.catalogString})"
             case _ => None

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullupCorrelatedPredicatesSuite.scala

@@ -0,0 +1,52 @@
+/*
+ * 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.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.{In, ListQuery}
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+
+class PullupCorrelatedPredicatesSuite extends PlanTest {
+
+  object Optimize extends RuleExecutor[LogicalPlan] {
+    val batches =
+      Batch("PullupCorrelatedPredicates", Once,
+        PullupCorrelatedPredicates) :: Nil
+  }
+
+  val testRelation = LocalRelation('a.int, 'b.double)
+  val testRelation2 = LocalRelation('c.int, 'd.double)
+
+  test("PullupCorrelatedPredicates should not produce unresolved plan") {
+    val correlatedSubquery =
+      testRelation2
+        .where('b < 'd)
+        .select('c)
+    val outerQuery =
+      testRelation
+        .where(In('a, Seq(ListQuery(correlatedSubquery))))
+        .select('a).analyze
+    assert(outerQuery.resolved)
+
+    val optimized = Optimize.execute(outerQuery)
+    assert(optimized.resolved)
+  }
+}