SPARK-23877: Fix edge case with unusual project/filter order.

Author: rdblue

sql/core/src/main/scala/org/apache/spark/sql/execution/OptimizeMetadataOnlyQuery.scala

@@ -49,9 +49,9 @@ case class OptimizeMetadataOnlyQuery(catalog: SessionCatalog) extends Rule[Logic
     }
 
     plan.transform {
-      case a @ Aggregate(_, aggExprs, child @ PartitionedRelation(partAttrs, filters, relation)) =>
+      case a @ Aggregate(_, aggExprs, child @ PartitionedRelation(_, attrs, filters, rel)) =>
         // We only apply this optimization when only partitioned attributes are scanned.
-        if (a.references.subsetOf(partAttrs)) {
+        if (a.references.subsetOf(attrs)) {
           val aggFunctions = aggExprs.flatMap(_.collect {
             case agg: AggregateExpression => agg
           })
@@ -67,7 +67,7 @@ case class OptimizeMetadataOnlyQuery(catalog: SessionCatalog) extends Rule[Logic
             })
           }
           if (isAllDistinctAgg) {
-            a.withNewChildren(Seq(replaceTableScanWithPartitionMetadata(child, relation, filters)))
+            a.withNewChildren(Seq(replaceTableScanWithPartitionMetadata(child, rel, filters)))
           } else {
             a
           }
@@ -159,26 +159,32 @@ case class OptimizeMetadataOnlyQuery(catalog: SessionCatalog) extends Rule[Logic
    */
   object PartitionedRelation extends PredicateHelper {
 
-    def unapply(plan: LogicalPlan): Option[(AttributeSet, Seq[Expression], LogicalPlan)] = {
+    def unapply(
+        plan: LogicalPlan): Option[(AttributeSet, AttributeSet, Seq[Expression], LogicalPlan)] = {
       plan match {
         case l @ LogicalRelation(fsRelation: HadoopFsRelation, _, _, _)
           if fsRelation.partitionSchema.nonEmpty =>
-          val partAttrs = getPartitionAttrs(fsRelation.partitionSchema.map(_.name), l)
-          Some((AttributeSet(partAttrs), Nil, l))
+          val partAttrs = AttributeSet(getPartitionAttrs(fsRelation.partitionSchema.map(_.name), l))
+          Some((partAttrs, partAttrs, Nil, l))
 
         case relation: HiveTableRelation if relation.tableMeta.partitionColumnNames.nonEmpty =>
-          val partAttrs = getPartitionAttrs(relation.tableMeta.partitionColumnNames, relation)
-          Some((AttributeSet(partAttrs), Nil, relation))
+          val partAttrs = AttributeSet(
+            getPartitionAttrs(relation.tableMeta.partitionColumnNames, relation))
+          Some((partAttrs, partAttrs, Nil, relation))
 
         case p @ Project(projectList, child) if projectList.forall(_.deterministic) =>
-          unapply(child).flatMap { case (partAttrs, filters, relation) =>
-            if (p.references.subsetOf(partAttrs)) Some((p.outputSet, filters, relation)) else None
+          unapply(child).flatMap { case (partAttrs, attrs, filters, relation) =>
+            if (p.references.subsetOf(attrs)) {
+              Some((partAttrs, p.outputSet, filters, relation))
+            } else {
+              None
+            }
           }
 
         case f @ Filter(condition, child) if condition.deterministic =>
-          unapply(child).flatMap { case (partAttrs, filters, relation) =>
+          unapply(child).flatMap { case (partAttrs, attrs, filters, relation) =>
             if (f.references.subsetOf(partAttrs)) {
-              Some((partAttrs, splitConjunctivePredicates(condition) ++ filters, relation))
+              Some((partAttrs, attrs, splitConjunctivePredicates(condition) ++ filters, relation))
             } else {
               None
             }

sql/hive/src/test/scala/org/apache/spark/sql/hive/OptimizeHiveMetadataOnlyQuerySuite.scala

@@ -21,22 +21,48 @@ import org.scalatest.BeforeAndAfter
 
 import org.apache.spark.metrics.source.HiveCatalogMetrics
 import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.catalyst.expressions.NamedExpression
+import org.apache.spark.sql.catalyst.plans.logical.{Distinct, Filter, Project, SubqueryAlias}
 import org.apache.spark.sql.hive.test.TestHiveSingleton
 import org.apache.spark.sql.test.SQLTestUtils
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
 
 class OptimizeHiveMetadataOnlyQuerySuite extends QueryTest with TestHiveSingleton
     with BeforeAndAfter with SQLTestUtils {
 
+  import spark.implicits._
+
+  before {
+    sql("CREATE TABLE metadata_only (id bigint, data string) PARTITIONED BY (part int)")
+    (0 to 10).foreach(p => sql(s"ALTER TABLE metadata_only ADD PARTITION (part=$p)"))
+  }
+
   test("SPARK-23877: validate metadata-only query pushes filters to metastore") {
     withTable("metadata_only") {
-      sql("CREATE TABLE metadata_only (id bigint, data string) PARTITIONED BY (part int)")
-      (0 to 100).foreach(p => sql(s"ALTER TABLE metadata_only ADD PARTITION (part=$p)"))
+      val startCount = HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount
 
       // verify the number of matching partitions
-      assert(sql("SELECT DISTINCT part FROM metadata_only WHERE part < 10").collect.size === 10)
+      assert(sql("SELECT DISTINCT part FROM metadata_only WHERE part < 5").collect().length === 5)
 
       // verify that the partition predicate was pushed down to the metastore
-      assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount === 10)
+      assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount - startCount === 5)
+    }
+  }
+
+  test("SPARK-23877: filter on projected expression") {
+    withTable("metadata_only") {
+      val startCount = HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount
+
+      // verify the matching partitions
+      val partitions = spark.internalCreateDataFrame(Distinct(Filter(($"x" < 5).expr,
+        Project(Seq(($"part" + 1).as("x").expr.asInstanceOf[NamedExpression]),
+          spark.table("metadata_only").logicalPlan.asInstanceOf[SubqueryAlias].child)))
+          .queryExecution.toRdd, StructType(Seq(StructField("x", IntegerType))))
+
+      checkAnswer(partitions, Seq(1, 2, 3, 4).toDF("x"))
+
+      // verify that the partition predicate was not pushed down to the metastore
+      assert(HiveCatalogMetrics.METRIC_PARTITIONS_FETCHED.getCount - startCount == 11)
     }
   }
 }