Merge pull request apache#48 from marmbrus/minorFixes

marmbrus · marmbrus · commit 222543100504 · 2014-02-24T15:18:21.000-08:00
Several minor fixes for bugs found during benchmarking.
diff --git a/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Catalog.scala b/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Catalog.scala
@@ -2,7 +2,7 @@ package org.apache.spark.sql
 package catalyst
 package analysis
 
-import plans.logical.LogicalPlan
+import plans.logical.{LogicalPlan, Subquery}
 import scala.collection.mutable
 
 /**
@@ -31,8 +31,14 @@ trait OverrideCatalog extends Catalog {
     tableName: String,
     alias: Option[String] = None): LogicalPlan = {
 
-    overrides.get((databaseName, tableName))
-      .getOrElse(super.lookupRelation(databaseName, tableName, alias))
+    val overriddenTable = overrides.get((databaseName, tableName))
+
+    // If an alias was specified by the lookup, wrap the plan in a subquery so that attributes are
+    // properly qualified with this alias.
+    val withAlias =
+      overriddenTable.map(r => alias.map(a => Subquery(a.toLowerCase, r)).getOrElse(r))
+
+    withAlias.getOrElse(super.lookupRelation(databaseName, tableName, alias))
   }
 
   def overrideTable(databaseName: Option[String], tableName: String, plan: LogicalPlan) =
diff --git a/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala b/catalyst/src/main/scala/org/apache/spark/sql/catalyst/dsl/package.scala
@@ -138,6 +138,9 @@ package object dsl {
         alias: Option[String] = None) =
       Generate(generator, join, outer, None, plan)
 
+    def insertInto(tableName: String, overwrite: Boolean = false) =
+      InsertIntoTable(analysis.UnresolvedRelation(None, tableName), Map.empty, plan, overwrite)
+
     def analyze = analysis.SimpleAnalyzer(plan)
   }
 }
diff --git a/catalyst/src/main/scala/org/apache/spark/sql/package.scala b/catalyst/src/main/scala/org/apache/spark/sql/package.scala
@@ -12,4 +12,6 @@ package object sql {
     com.typesafe.scalalogging.slf4j.Logger(org.slf4j.LoggerFactory.getLogger(name))
 
   protected[sql] type Logging = com.typesafe.scalalogging.slf4j.Logging
+
+  type Row = catalyst.expressions.Row
 }
diff --git a/core/src/main/scala/org/apache/spark/rdd/SharkPairRDDFunctions.scala b/core/src/main/scala/org/apache/spark/rdd/SharkPairRDDFunctions.scala
@@ -56,7 +56,7 @@ class SharkPairRDDFunctions[K: ClassTag, V: ClassTag](self: RDD[(K, V)])
   def groupByKeyLocally(): RDD[(K, Seq[V])] = {
     def createCombiner(v: V) = ArrayBuffer(v)
     def mergeValue(buf: ArrayBuffer[V], v: V) = buf += v
-    val aggregator = new Aggregator[K, V, ArrayBuffer[V]](createCombiner _, mergeValue _, null)
+    val aggregator = new Aggregator[K, V, ArrayBuffer[V]](createCombiner _, mergeValue _, _ ++ _)
     val bufs = self.mapPartitionsWithContext((context, iter) => {
       new InterruptibleIterator(context, aggregator.combineValuesByKey(iter, context))
     }, preservesPartitioning = true)
diff --git a/core/src/main/scala/org/apache/spark/sql/SparkSqlContext.scala b/core/src/main/scala/org/apache/spark/sql/SparkSqlContext.scala
@@ -60,6 +60,7 @@ class SparkSqlContext(val sparkContext: SparkContext) extends Logging {
     val sparkContext = self.sparkContext
 
     val strategies: Seq[Strategy] =
+      TopK ::
       PartialAggregation ::
       SparkEquiInnerJoin ::
       BasicOperators ::
diff --git a/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -137,8 +137,25 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
   protected lazy val singleRowRdd =
     sparkContext.parallelize(Seq(new GenericRow(IndexedSeq()): Row), 1)
 
+  def convertToCatalyst(a: Any): Any = a match {
+    case s: Seq[Any] => s.map(convertToCatalyst)
+    case p: Product => new GenericRow(p.productIterator.map(convertToCatalyst).toSeq)
+    case other => other
+  }
+
+  object TopK extends Strategy {
+    def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
+      case logical.StopAfter(limit, logical.Sort(order, child)) =>
+        execution.TopK(
+          Evaluate(limit, Nil).asInstanceOf[Int], order, planLater(child))(sparkContext) :: Nil
+      case _ => Nil
+    }
+  }
+
   // Can we automate these 'pass through' operations?
   object BasicOperators extends Strategy {
+    // TOOD: Set
+    val numPartitions = 200
     def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
       case logical.Distinct(child) =>
         execution.Aggregate(
@@ -160,7 +177,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.Sample(fraction, withReplacement, seed, planLater(child)) :: Nil
       case logical.LocalRelation(output, data) =>
         val dataAsRdd =
-          sparkContext.parallelize(data.map(r => new GenericRow(r.productIterator.toVector): Row))
+          sparkContext.parallelize(data.map(r =>
+            new GenericRow(r.productIterator.map(convertToCatalyst).toVector): Row))
         execution.ExistingRdd(output, dataAsRdd) :: Nil
       case logical.StopAfter(limit, child) =>
         execution.StopAfter(
@@ -172,6 +190,8 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         execution.Generate(generator, join = join, outer = outer, planLater(child)) :: Nil
       case logical.NoRelation =>
         execution.ExistingRdd(Nil, singleRowRdd) :: Nil
+      case logical.Repartition(expressions, child) =>
+        execution.Exchange(HashPartitioning(expressions, numPartitions), planLater(child)) :: Nil
       case _ => Nil
     }
   }
diff --git a/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala b/core/src/main/scala/org/apache/spark/sql/execution/basicOperators.scala
@@ -53,6 +53,23 @@ case class StopAfter(limit: Int, child: SparkPlan)(@transient sc: SparkContext)
   def execute() = sc.makeRDD(executeCollect(), 1)
 }
 
+case class TopK(limit: Int, sortOrder: Seq[SortOrder], child: SparkPlan)
+               (@transient sc: SparkContext) extends UnaryNode {
+  override def otherCopyArgs = sc :: Nil
+
+  def output = child.output
+
+  @transient
+  lazy val ordering = new RowOrdering(sortOrder)
+
+  override def executeCollect() = child.execute().takeOrdered(limit)(ordering)
+
+  // TODO: Terminal split should be implemented differently from non-terminal split.
+  // TODO: Pick num splits based on |limit|.
+  def execute() = sc.makeRDD(executeCollect(), 1)
+}
+
+
 case class Sort(
     sortOrder: Seq[SortOrder],
     global: Boolean,
diff --git a/core/src/main/scala/org/apache/spark/sql/execution/package.scala b/core/src/main/scala/org/apache/spark/sql/execution/package.scala
@@ -4,5 +4,4 @@ package org.apache.spark.sql
  * An execution engine for relational query plans that runs on top Spark and returns RDDs.
  */
 package object execution {
-  type Row = catalyst.expressions.Row
 }
diff --git a/core/src/test/scala/org/apache/spark/sql/DslQueryTests.scala b/core/src/test/scala/org/apache/spark/sql/DslQueryTests.scala
@@ -96,6 +96,13 @@ class BasicQuerySuite extends DslQueryTest {
       testData.data)
   }
 
+  test("agg") {
+    checkAnswer(
+      testData2.groupBy('a)('a, Sum('b)),
+      Seq((1,3),(2,3),(3,3))
+    )
+  }
+
   test("select *") {
     checkAnswer(
       testData.select(Star(None)),
diff --git a/shark/src/main/scala/org/apache/spark/sql/shark/HiveMetastoreCatalog.scala b/shark/src/main/scala/org/apache/spark/sql/shark/HiveMetastoreCatalog.scala
@@ -20,6 +20,8 @@ import catalyst.types._
 import scala.collection.JavaConversions._
 
 class HiveMetastoreCatalog(shark: SharkContext) extends Catalog with Logging {
+  import HiveMetastoreTypes._
+
   val client = Hive.get(shark.hiveconf)
 
   def lookupRelation(
@@ -42,37 +44,39 @@ class HiveMetastoreCatalog(shark: SharkContext) extends Catalog with Logging {
       alias)(table.getTTable, partitions.map(part => part.getTPartition))
   }
 
+  def createTable(databaseName: String, tableName: String, schema: Seq[Attribute]) {
+    val table = new Table(databaseName, tableName)
+    val hiveSchema =
+      schema.map(attr => new FieldSchema(attr.name, toMetastoreType(attr.dataType), ""))
+    table.setFields(hiveSchema)
+
+    val sd = new StorageDescriptor()
+    table.getTTable.setSd(sd)
+    sd.setCols(hiveSchema)
+
+    // TODO: THESE ARE ALL DEFAULTS, WE NEED TO PARSE / UNDERSTAND the output specs.
+    sd.setCompressed(false)
+    sd.setParameters(Map[String, String]())
+    sd.setInputFormat("org.apache.hadoop.mapred.TextInputFormat")
+    sd.setOutputFormat("org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat")
+    val serDeInfo = new SerDeInfo()
+    serDeInfo.setName(tableName)
+    serDeInfo.setSerializationLib("org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe")
+    serDeInfo.setParameters(Map[String, String]())
+    sd.setSerdeInfo(serDeInfo)
+    client.createTable(table)
+  }
+
   /**
    * Creates any tables required for query execution.
    * For example, because of a CREATE TABLE X AS statement.
    */
   object CreateTables extends Rule[LogicalPlan] {
-    import HiveMetastoreTypes._
-
     def apply(plan: LogicalPlan): LogicalPlan = plan transform {
       case InsertIntoCreatedTable(db, tableName, child) =>
         val databaseName = db.getOrElse(SessionState.get.getCurrentDatabase())
 
-        val table = new Table(databaseName, tableName)
-        val schema =
-          child.output.map(attr => new FieldSchema(attr.name, toMetastoreType(attr.dataType), ""))
-        table.setFields(schema)
-
-        val sd = new StorageDescriptor()
-        table.getTTable.setSd(sd)
-        sd.setCols(schema)
-
-        // TODO: THESE ARE ALL DEFAULTS, WE NEED TO PARSE / UNDERSTAND the output specs.
-        sd.setCompressed(false)
-        sd.setParameters(Map[String, String]())
-        sd.setInputFormat("org.apache.hadoop.mapred.TextInputFormat")
-        sd.setOutputFormat("org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat")
-        val serDeInfo = new SerDeInfo()
-        serDeInfo.setName(tableName)
-        serDeInfo.setSerializationLib("org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe")
-        serDeInfo.setParameters(Map[String, String]())
-        sd.setSerdeInfo(serDeInfo)
-        client.createTable(table)
+        createTable(databaseName, tableName, child.output)
 
         InsertIntoTable(
           lookupRelation(Some(databaseName), tableName, None).asInstanceOf[BaseRelation],
diff --git a/shark/src/main/scala/org/apache/spark/sql/shark/SharkContext.scala b/shark/src/main/scala/org/apache/spark/sql/shark/SharkContext.scala
@@ -26,8 +26,11 @@ import scala.collection.JavaConversions._
  * Starts up an instance of shark where metadata is stored locally. An in-process metadata data is
  * created with data stored in ./metadata.  Warehouse data is stored in in ./warehouse.
  */
-class LocalSharkContext(sc: SparkContext) extends SharkContext(sc) {
-  override def warehousePath = new File("warehouse").getCanonicalPath
+class LocalSharkContext(
+    sc: SparkContext,
+    override val warehousePath: String = new File("warehouse").getCanonicalPath)
+  extends SharkContext(sc) {
+
   override def metastorePath = new File("metastore").getCanonicalPath
 }
 
@@ -172,6 +175,7 @@ abstract class SharkContext(sc: SparkContext) extends SparkSqlContext(sc) {
     val sharkContext = self
 
     override val strategies: Seq[Strategy] = Seq(
+      TopK,
       ColumnPrunings,
       PartitionPrunings,
       HiveTableScans,
diff --git a/shark/src/main/scala/org/apache/spark/sql/shark/hiveOperators.scala b/shark/src/main/scala/org/apache/spark/sql/shark/hiveOperators.scala
@@ -180,8 +180,16 @@ case class InsertIntoHiveTable(
     case (s: String, oi: JavaHiveVarcharObjectInspector) => new HiveVarchar(s, s.size)
     case (bd: BigDecimal, oi: JavaHiveDecimalObjectInspector) =>
       new HiveDecimal(bd.underlying())
+    case (row: Row, oi: StandardStructObjectInspector) =>
+      val struct = oi.create()
+      row.zip(oi.getAllStructFieldRefs).foreach {
+        case (data, field) =>
+          oi.setStructFieldData(struct, field, wrap(data, field.getFieldObjectInspector))
+      }
+      struct
     case (s: Seq[_], oi: ListObjectInspector) =>
-      seqAsJavaList(s.map(wrap(_, oi.getListElementObjectInspector)))
+      val wrappedSeq = s.map(wrap(_, oi.getListElementObjectInspector))
+      seqAsJavaList(wrappedSeq)
     case (obj, _) => obj
   }
 

Original file line number	Diff line number	Diff line change
`@@ -138,6 +138,9 @@ package object dsl {`
`138`	`138`	`alias: Option[String] = None) =`
`139`	`139`	`Generate(generator, join, outer, None, plan)`
`140`	`140`
	`141`	`+ def insertInto(tableName: String, overwrite: Boolean = false) =`
	`142`	`+ InsertIntoTable(analysis.UnresolvedRelation(None, tableName), Map.empty, plan, overwrite)`
	`143`	`+`
`141`	`144`	`def analyze = analysis.SimpleAnalyzer(plan)`
`142`	`145`	`}`
`143`	`146`	`}`
Original file line number	Diff line number	Diff line change
`@@ -12,4 +12,6 @@ package object sql {`
`12`	`12`	`com.typesafe.scalalogging.slf4j.Logger(org.slf4j.LoggerFactory.getLogger(name))`
`13`	`13`
`14`	`14`	`protected[sql] type Logging = com.typesafe.scalalogging.slf4j.Logging`
	`15`	`+`
	`16`	`+ type Row = catalyst.expressions.Row`
`15`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -4,5 +4,4 @@ package org.apache.spark.sql`
`4`	`4`	`* An execution engine for relational query plans that runs on top Spark and returns RDDs.`
`5`	`5`	`*/`
`6`	`6`	`package object execution {`
`7`		`- type Row = catalyst.expressions.Row`
`8`	`7`	`}`