update

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala

@@ -27,9 +27,11 @@ import org.apache.spark.sql.{AnalysisException, SaveMode, SparkSession}
 import org.apache.spark.sql.catalyst.{InternalRow, TableIdentifier}
 import org.apache.spark.sql.catalyst.catalog._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.logical
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules._
+import org.apache.spark.sql.execution.LogicalRDD
 import org.apache.spark.sql.execution.command.CreateDataSourceTableUtils._
 import org.apache.spark.sql.execution.command.CreateHiveTableAsSelectLogicalPlan
 import org.apache.spark.sql.execution.datasources.{Partition => _, _}
@@ -506,6 +508,119 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log
     }
   }
 
+  /**
+   * When scanning only partition columns, get results based on metadata without scanning files.
+   * It is used for distinct or distinct/Max/Min aggregations, example: max(partition).
+   */
+  object MetadataOnlyOptimizer extends Rule[LogicalPlan] {
+
+    private def canSupportMetadataOnly(a: Aggregate): Boolean = {
+      val aggregateExpressions = a.aggregateExpressions.flatMap { expr =>
+        expr.collect {
+          case agg: AggregateExpression => agg
+        }
+      }.distinct
+      aggregateExpressions.forall { agg =>
+        if (agg.isDistinct) {
+          true
+        } else {
+          agg.aggregateFunction match {
+            case max: Max => true
+            case min: Min => true
+            case _ => false
+          }
+        }
+      }
+    }
+
+    private def findRelation(plan: LogicalPlan): (Option[LogicalPlan], Seq[Expression]) = {
+      plan match {
+        case relation @ LogicalRelation(files: HadoopFsRelation, _, table)
+          if files.partitionSchema.nonEmpty =>
+          (Some(relation), Seq.empty[Expression])
+
+        case relation: MetastoreRelation if relation.partitionKeys.nonEmpty =>
+          (Some(relation), Seq.empty[Expression])
+
+        case p @ Project(_, child) =>
+          findRelation(child)
+
+        case f @ Filter(filterCondition, child) =>
+          val (plan, conditions) = findRelation(child)
+          (plan, conditions ++ Seq(filterCondition))
+
+        case SubqueryAlias(_, child) =>
+          findRelation(child)
+
+        case _ => (None, Seq.empty[Expression])
+      }
+    }
+
+    private def convertToMetadataOnlyPlan(
+        parent: LogicalPlan,
+        filters: Seq[Expression],
+        relation: LogicalPlan): LogicalPlan = relation match {
+      case l @ LogicalRelation(files: HadoopFsRelation, _, _) =>
+        val attributeMap = l.output.map(attr => (attr.name, attr)).toMap
+        val partitionColumns = files.partitionSchema.map { field =>
+          attributeMap.getOrElse(field.name, throw new AnalysisException(
+            s"Unable to resolve ${field.name} given [${l.output.map(_.name).mkString(", ")}]"))
+        }
+        val filterColumns = filters.flatMap(_.references)
+        val projectSet = parent.references ++ AttributeSet(filterColumns)
+        if (projectSet.subsetOf(AttributeSet(partitionColumns))) {
+          val selectedPartitions = files.location.listFiles(filters)
+          val partitionValues = selectedPartitions.map(_.values)
+          val valuesRdd = sparkSession.sparkContext.parallelize(partitionValues, 1)
+          parent.withNewChildren(LogicalRDD(partitionColumns, valuesRdd)(sparkSession) :: Nil)
+        } else {
+          parent
+        }
+
+      case relation: MetastoreRelation =>
+        if (parent.references.subsetOf(AttributeSet(relation.partitionKeys))) {
+          val partitionColumnDataTypes = relation.partitionKeys.map(_.dataType)
+          val partitionValues = relation.getHiveQlPartitions(filters).map { p =>
+            InternalRow.fromSeq(p.getValues.asScala.zip(partitionColumnDataTypes).map {
+              case (rawValue, dataType) => Cast(Literal(rawValue), dataType).eval(null)
+            })
+          }
+          val valuesRdd = sparkSession.sparkContext.parallelize(partitionValues, 1)
+          val valuesPlan = LogicalRDD(relation.partitionKeys, valuesRdd)(sparkSession)
+          val child = filters.reduceLeftOption(And).map(Filter(_, valuesPlan)).getOrElse(valuesPlan)
+          parent.withNewChildren(child :: Nil)
+        } else {
+          parent
+        }
+
+      case _ =>
+        parent
+    }
+
+    def apply(plan: LogicalPlan): LogicalPlan = {
+      if (!sparkSession.sessionState.conf.optimizerMetadataOnly) {
+        return plan
+      }
+      plan.transform {
+        case a @ Aggregate(_, _, child) if canSupportMetadataOnly(a) =>
+          val (plan, filters) = findRelation(child)
+          if (plan.isDefined) {
+            convertToMetadataOnlyPlan(a, filters, plan.get)
+          } else {
+            a
+          }
+
+        case d @ Distinct(p @ Project(_, _)) =>
+          val (plan, filters) = findRelation(p)
+          if (plan.isDefined) {
+            d.withNewChildren(convertToMetadataOnlyPlan(p, filters, plan.get) :: Nil)
+          } else {
+            d
+          }
+      }
+    }
+  }
+
 }
 
 /**

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionCatalog.scala

@@ -88,6 +88,7 @@ private[sql] class HiveSessionCatalog(
   val OrcConversions: Rule[LogicalPlan] = metastoreCatalog.OrcConversions
   val CreateTables: Rule[LogicalPlan] = metastoreCatalog.CreateTables
   val PreInsertionCasts: Rule[LogicalPlan] = metastoreCatalog.PreInsertionCasts
+  val metadataOnlyOptimizer: Rule[LogicalPlan] = metastoreCatalog.MetadataOnlyOptimizer
 
   override def refreshTable(name: TableIdentifier): Unit = {
     metastoreCatalog.refreshTable(name)

sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionState.scala

@@ -19,7 +19,8 @@ package org.apache.spark.sql.hive
 
 import org.apache.spark.sql._
 import org.apache.spark.sql.catalyst.analysis.Analyzer
-import org.apache.spark.sql.execution.SparkPlanner
+import org.apache.spark.sql.catalyst.optimizer.Optimizer
+import org.apache.spark.sql.execution.{SparkOptimizer, SparkPlanner}
 import org.apache.spark.sql.execution.datasources._
 import org.apache.spark.sql.hive.client.HiveClient
 import org.apache.spark.sql.internal.SessionState
@@ -74,6 +75,18 @@ private[hive] class HiveSessionState(sparkSession: SparkSession)
     }
   }
 
+  /**
+   * Logical query plan optimizer for Hive.
+   */
+  override lazy val optimizer: Optimizer = new SparkOptimizer(catalog, conf, experimentalMethods) {
+    val extendedHiveOptimizerRules: Seq[Batch] = Seq(
+      Batch("MetadataOnly Optimization", Once,
+        catalog.metadataOnlyOptimizer)
+    )
+
+    override def batches: Seq[Batch] = extendedHiveOptimizerRules ++ super.batches
+  }
+
   /**
    * Planner that takes into account Hive-specific strategies.
    */

sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala

@@ -1640,4 +1640,44 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
       )
     }
   }
+
+  test("spark-15752 metadata only optimizer") {
+    withSQLConf(SQLConf.OPTIMIZER_METADATA_ONLY.key -> "true") {
+      spark.range(0, 5, 1, 2).selectExpr("id as key", "id as value").registerTempTable("tempTable")
+      sql(
+        """
+          |CREATE TABLE srcpart_15752 (key INT, value STRING)
+          |PARTITIONED BY (ds STRING, hr INT) STORED AS parquet
+        """.stripMargin)
+      for (ds <- Seq("2008-04-08", "2008-04-09"); hr <- Seq(11, 12)) {
+        sql(
+          s"""
+             |INSERT OVERWRITE TABLE srcpart_15752 PARTITION (ds='$ds',hr='$hr')
+             |select key, value from tempTable
+        """.stripMargin)
+      }
+      checkAnswer(sql("select max(hr) from srcpart_15752"), Row(12))
+      checkAnswer(sql("select max(hr) from srcpart_15752 where hr = 11"), Row(11))
+      checkAnswer(sql("select max(hr) from (select hr from srcpart_15752) t"), Row(12))
+      checkAnswer(sql("select distinct hr from srcpart_15752 where hr = 11"), Row(11))
+
+      sql(
+        """
+          |CREATE TABLE srctext_15752 (key INT, value STRING)
+          |PARTITIONED BY (ds STRING, hr INT) STORED AS textfile
+        """.stripMargin)
+      for (ds <- Seq("2008-04-08", "2008-04-09"); hr <- Seq(11, 12)) {
+        sql(
+          s"""
+             |INSERT OVERWRITE TABLE srctext_15752 PARTITION (ds='$ds',hr='$hr')
+             |select key, value from tempTable
+        """.stripMargin)
+      }
+
+      checkAnswer(sql("select max(hr) from srctext_15752"), Row(12))
+      checkAnswer(sql("select max(hr) from srctext_15752 where hr = 11"), Row(11))
+      checkAnswer(sql("select max(hr) from (select hr from srctext_15752) t"), Row(12))
+      checkAnswer(sql("select distinct hr from srctext_15752 where hr = 11"), Row(11))
+    }
+  }
 }