Merge branch 'apache-master' into metadata-only

lianhuiwang · lianhuiwang · commit 2e55a9df94ef · 2016-06-24T15:50:56.000+08:00
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/ExistingRDD.scala
@@ -347,8 +347,7 @@ private[sql] object DataSourceScanExec {
       rdd: RDD[InternalRow],
       relation: BaseRelation,
       metadata: Map[String, String] = Map.empty,
-      metastoreTableIdentifier: Option[TableIdentifier] = None,
-      isSupportBatch: Boolean = true): DataSourceScanExec = {
+      metastoreTableIdentifier: Option[TableIdentifier] = None): DataSourceScanExec = {
     val outputPartitioning = {
       val bucketSpec = relation match {
         // TODO: this should be closer to bucket planning.
@@ -372,8 +371,7 @@ private[sql] object DataSourceScanExec {
 
     relation match {
       case r: HadoopFsRelation
-        if isSupportBatch &&
-          r.fileFormat.supportBatch(r.sparkSession, StructType.fromAttributes(output)) =>
+        if r.fileFormat.supportBatch(r.sparkSession, StructType.fromAttributes(output)) =>
         BatchedDataSourceScanExec(
           output, rdd, relation, outputPartitioning, metadata, metastoreTableIdentifier)
       case _ =>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/FileSourceStrategySuite.scala
@@ -338,21 +338,6 @@ class FileSourceStrategySuite extends QueryTest with SharedSQLContext with Predi
         assert(partitions.flatMap(fileScanRDD.preferredLocations).length == 3)
       }
     }
-
-    test("optimize metadataOnly") {
-      withSQLConf("spark.sql.optimizer.metadataOnly" -> "true") {
-        val table =
-          createTable(
-            files = Seq(
-              "p1=1/file1" -> 10,
-              "p1=2/file2" -> 10))
-
-        checkDataset(table.select($"p1"), Row(1), Row(2))
-        checkDataset(table.where("p1 = 1").select($"p1"), Row(1))
-        val df = table.where("p1 = 1 AND (p1 + c1) = 2 AND c1 = 1")
-        assert(getPhysicalFilters(df) contains resolve(df, "c1 = 1"))
-      }
-    }
   }
 
   test("SPARK-15654 do not split non-splittable files") {
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveMetastoreCatalog.scala
@@ -23,21 +23,22 @@ import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache}
 import org.apache.hadoop.fs.Path
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.sql.{AnalysisException, SaveMode, SparkSession}
+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
+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.
  *
@@ -355,13 +356,7 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log
       val fileFormatClass = classOf[ParquetFileFormat]
 
       val mergeSchema = sessionState.convertMetastoreParquetWithSchemaMerging
-      val options = Map(
-        ParquetFileFormat.MERGE_SCHEMA -> mergeSchema.toString,
-        ParquetFileFormat.METASTORE_TABLE_NAME -> TableIdentifier(
-          relation.tableName,
-          Some(relation.databaseName)
-        ).unquotedString
-      )
+      val options = Map(ParquetOptions.MERGE_SCHEMA -> mergeSchema.toString)
 
       convertToLogicalRelation(relation, options, defaultSource, fileFormatClass, "parquet")
     }
@@ -465,43 +460,119 @@ private[hive] class HiveMetastoreCatalog(sparkSession: SparkSession) extends Log
   }
 
   /**
-   * Casts input data to correct data types according to table definition before inserting into
-   * that table.
+   * 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 PreInsertionCasts extends Rule[LogicalPlan] {
-    def apply(plan: LogicalPlan): LogicalPlan = plan.transform {
-      // Wait until children are resolved.
-      case p: LogicalPlan if !p.childrenResolved => p
+  object MetadataOnlyOptimizer extends Rule[LogicalPlan] {
 
-      case p @ InsertIntoTable(table: MetastoreRelation, _, child, _, _) =>
-        castChildOutput(p, table, child)
+    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
+          }
+        }
+      }
     }
 
-    def castChildOutput(p: InsertIntoTable, table: MetastoreRelation, child: LogicalPlan)
-      : LogicalPlan = {
-      val childOutputDataTypes = child.output.map(_.dataType)
-      val numDynamicPartitions = p.partition.values.count(_.isEmpty)
-      val tableOutputDataTypes =
-        (table.attributes ++ table.partitionKeys.takeRight(numDynamicPartitions))
-          .take(child.output.length).map(_.dataType)
-
-      if (childOutputDataTypes == tableOutputDataTypes) {
-        InsertIntoHiveTable(table, p.partition, p.child, p.overwrite, p.ifNotExists)
-      } else if (childOutputDataTypes.size == tableOutputDataTypes.size &&
-        childOutputDataTypes.zip(tableOutputDataTypes)
-          .forall { case (left, right) => left.sameType(right) }) {
-        // If both types ignoring nullability of ArrayType, MapType, StructType are the same,
-        // use InsertIntoHiveTable instead of InsertIntoTable.
-        InsertIntoHiveTable(table, p.partition, p.child, p.overwrite, p.ifNotExists)
-      } else {
-        // Only do the casting when child output data types differ from table output data types.
-        val castedChildOutput = child.output.zip(table.output).map {
-          case (input, output) if input.dataType != output.dataType =>
-            Alias(Cast(input, output.dataType), input.name)()
-          case (input, _) => input
+    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
         }
 
-        p.copy(child = logical.Project(castedChildOutput, child))
+      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
+          }
       }
     }
   }
diff --git a/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionCatalog.scala b/sql/hive/src/main/scala/org/apache/spark/sql/hive/HiveSessionCatalog.scala
@@ -87,7 +87,6 @@ private[sql] class HiveSessionCatalog(
   val ParquetConversions: Rule[LogicalPlan] = metastoreCatalog.ParquetConversions
   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 = {
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/SQLQuerySuite.scala
@@ -1687,7 +1687,8 @@ 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")
+      val df = Seq((1, 2), (3, 4)).toDF("key", "value")
+      df.createOrReplaceTempView("data")
       sql(
         """
           |CREATE TABLE srcpart_15752 (key INT, value STRING)
@@ -1697,12 +1698,13 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
         sql(
           s"""
              |INSERT OVERWRITE TABLE srcpart_15752 PARTITION (ds='$ds',hr='$hr')
-             |select key, value from tempTable
+             |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(
@@ -1714,13 +1716,14 @@ class SQLQuerySuite extends QueryTest with SQLTestUtils with TestHiveSingleton {
         sql(
           s"""
              |INSERT OVERWRITE TABLE srctext_15752 PARTITION (ds='$ds',hr='$hr')
-             |select key, value from tempTable
+             |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))
     }
   }
