[SPARK-15752] [SQL] Optimize metadata only query that has an aggregate whose children are deterministic project or filter operators.

sql/core/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -258,6 +258,11 @@ object SQLConf {
       .booleanConf
       .createWithDefault(false)
 
+  val OPTIMIZER_METADATA_ONLY = SQLConfigBuilder("spark.sql.optimizer.metadataOnly")
+    .doc("When true, enable the metadata-only query optimization.")
+    .booleanConf
+    .createWithDefault(false)
+
   val NATIVE_VIEW = SQLConfigBuilder("spark.sql.nativeView")
     .internal()
     .doc("When true, CREATE VIEW will be handled by Spark SQL instead of Hive native commands.  " +
@@ -613,6 +618,8 @@ private[sql] class SQLConf extends Serializable with CatalystConf with Logging {
 
   def metastorePartitionPruning: Boolean = getConf(HIVE_METASTORE_PARTITION_PRUNING)
 
+  def optimizerMetadataOnly: Boolean = getConf(OPTIMIZER_METADATA_ONLY)
+
   def nativeView: Boolean = getConf(NATIVE_VIEW)
 
   def wholeStageEnabled: Boolean = getConf(WHOLESTAGE_CODEGEN_ENABLED)

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

@@ -27,17 +27,18 @@ import org.apache.spark.sql.{AnalysisException, 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 => _, _}
 import org.apache.spark.sql.execution.datasources.parquet.{ParquetFileFormat, ParquetOptions}
 import org.apache.spark.sql.hive.orc.OrcFileFormat
 import org.apache.spark.sql.types._
 
-
 /**
  * Legacy catalog for interacting with the Hive metastore.
  *
@@ -457,6 +458,125 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log
           allowExisting)
     }
   }
+
+  /**
+   * 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,
+        project: Option[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)
+          val valuesPlan = LogicalRDD(partitionColumns, valuesRdd)(sparkSession)
+          val scanPlan = project.map(_.withNewChildren(valuesPlan :: Nil)).getOrElse(valuesPlan)
+          parent.withNewChildren(scanPlan :: 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 filterPlan =
+            filters.reduceLeftOption(And).map(Filter(_, valuesPlan)).getOrElse(valuesPlan)
+          val scanPlan = project.map(_.withNewChildren(filterPlan :: Nil)).getOrElse(filterPlan)
+          parent.withNewChildren(scanPlan :: 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, None, filters, plan.get)
+          } else {
+            a
+          }
+
+        case d @ Distinct(p @ Project(_, _)) =>
+          val (plan, filters) = findRelation(p)
+          if (plan.isDefined) {
+            convertToMetadataOnlyPlan(d, Some(p), filters, plan.get)
+          } else {
+            d
+          }
+      }
+    }
+  }
+
 }
 
 /**

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

@@ -87,6 +87,7 @@ private[sql] class HiveSessionCatalog(
   val ParquetConversions: Rule[LogicalPlan] = metastoreCatalog.ParquetConversions
   val OrcConversions: Rule[LogicalPlan] = metastoreCatalog.OrcConversions
   val CreateTables: Rule[LogicalPlan] = metastoreCatalog.CreateTables
+  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
@@ -73,6 +74,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

@@ -1684,4 +1684,47 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
       )
     }
   }
+
+  test("spark-15752 metadata only optimizer") {
+    withSQLConf(SQLConf.OPTIMIZER_METADATA_ONLY.key -> "true") {
+      val df = Seq((1, 2), (3, 4)).toDF("key", "value")
+      df.createOrReplaceTempView("data")
+      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 data
+        """.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"), Row(11) :: Row(12) :: Nil)
+      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 data
+        """.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"), Row(11) :: Row(12) :: Nil)
+      checkAnswer(sql("select distinct hr from srctext_15752 where hr = 11"), Row(11))
+    }
+  }
 }