Added "(UN)CACHE TABLE" SQL/HiveQL statements

Author: liancheng

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/SqlParser.scala

@@ -111,6 +111,7 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
   protected val APPROXIMATE = Keyword("APPROXIMATE")
   protected val AVG = Keyword("AVG")
   protected val BY = Keyword("BY")
+  protected val CACHE = Keyword("CACHE")
   protected val CAST = Keyword("CAST")
   protected val COUNT = Keyword("COUNT")
   protected val DESC = Keyword("DESC")
@@ -149,7 +150,9 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
   protected val SEMI = Keyword("SEMI")
   protected val STRING = Keyword("STRING")
   protected val SUM = Keyword("SUM")
+  protected val TABLE = Keyword("TABLE")
   protected val TRUE = Keyword("TRUE")
+  protected val UNCACHE = Keyword("UNCACHE")
   protected val UNION = Keyword("UNION")
   protected val WHERE = Keyword("WHERE")
 
@@ -189,7 +192,7 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
         UNION ~ ALL ^^^ { (q1: LogicalPlan, q2: LogicalPlan) => Union(q1, q2) } |
         UNION ~ opt(DISTINCT) ^^^ { (q1: LogicalPlan, q2: LogicalPlan) => Distinct(Union(q1, q2)) }
       )
-    | insert
+    | insert | cache
   )
 
   protected lazy val select: Parser[LogicalPlan] =
@@ -220,6 +223,11 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
         InsertIntoTable(r, Map[String, Option[String]](), s, overwrite)
     }
 
+  protected lazy val cache: Parser[LogicalPlan] =
+    (CACHE ^^^ true | UNCACHE ^^^ false) ~ TABLE ~ ident ^^ {
+      case doCache ~ _ ~ tableName => CacheCommand(tableName, doCache)
+    }
+
   protected lazy val projections: Parser[Seq[Expression]] = repsep(projection, ",")
 
   protected lazy val projection: Parser[Expression] =

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala

@@ -129,6 +129,11 @@ case class ExplainCommand(plan: LogicalPlan) extends Command {
   override def output = Seq(AttributeReference("plan", StringType, nullable = false)())
 }
 
+/**
+ * Returned for the "CACHE TABLE tableName" and "UNCACHE TABLE tableName" command.
+ */
+case class CacheCommand(tableName: String, doCache: Boolean) extends Command
+
 /**
  * A logical plan node with single child.
  */

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

@@ -239,10 +239,11 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
       case logical.SetCommand(key, value) =>
         Seq(execution.SetCommandPhysical(key, value, plan.output)(context))
       case logical.ExplainCommand(child) =>
-        val qe = context.executePlan(child)
-        Seq(execution.ExplainCommandPhysical(qe.executedPlan, plan.output)(context))
+        val executedPlan = context.executePlan(child).executedPlan
+        Seq(execution.ExplainCommandPhysical(executedPlan, plan.output)(context))
+      case logical.CacheCommand(tableName, cache) =>
+        Seq(execution.CacheCommandPhysical(tableName, cache)(context))
       case _ => Nil
     }
   }
-
 }

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

@@ -65,3 +65,26 @@ case class ExplainCommandPhysical(child: SparkPlan, output: Seq[Attribute])
 
   override def otherCopyArgs = context :: Nil
 }
+
+/**
+ * :: DeveloperApi ::
+ */
+@DeveloperApi
+case class CacheCommandPhysical(tableName: String, doCache: Boolean)(@transient context: SQLContext)
+  extends LeafNode {
+
+  lazy val commandSideEffect = {
+    if (doCache) {
+      context.cacheTable(tableName)
+    } else {
+      context.uncacheTable(tableName)
+    }
+  }
+
+  override def execute(): RDD[Row] = {
+    commandSideEffect
+    context.emptyResult
+  }
+
+  override def output: Seq[Attribute] = Seq.empty
+}

sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala

@@ -70,4 +70,18 @@ class CachedTableSuite extends QueryTest {
     TestSQLContext.sql("SELECT * FROM testData a JOIN testData b ON a.key = b.key")
     TestSQLContext.uncacheTable("testData")
   }
+
+  test("'CACHE TABLE' and 'UNCACHE TABLE' SQL statement") {
+    TestSQLContext.sql("CACHE TABLE testData")
+    TestSQLContext.table("testData").queryExecution.executedPlan match {
+      case _: InMemoryColumnarTableScan => // Found evidence of caching
+      case _ => fail(s"Table 'testData' should be cached")
+    }
+
+    TestSQLContext.sql("UNCACHE TABLE testData")
+    TestSQLContext.table("testData").queryExecution.executedPlan match {
+      case _: InMemoryColumnarTableScan => fail(s"Table 'testData' should not be cached")
+      case _ => // Found evidence of uncaching
+    }
+  }
 }

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

@@ -218,6 +218,10 @@ private[hive] object HiveQl {
           case Array(key, value) => // "set key=value"
             SetCommand(Some(key), Some(value))
         }
+      } else if (sql.toLowerCase.startsWith("cache table")) {
+        CacheCommand(sql.drop(12).trim, true)
+      } else if (sql.toLowerCase.startsWith("uncache table")) {
+        CacheCommand(sql.drop(14).trim, false)
       } else if (sql.toLowerCase.startsWith("add jar")) {
         AddJar(sql.drop(8))
       } else if (sql.toLowerCase.startsWith("add file")) {
@@ -839,11 +843,11 @@ private[hive] object HiveQl {
     case Token("TOK_FUNCTIONDI", Token(SUM(), Nil) :: arg :: Nil) => SumDistinct(nodeToExpr(arg))
     case Token("TOK_FUNCTION", Token(MAX(), Nil) :: arg :: Nil) => Max(nodeToExpr(arg))
     case Token("TOK_FUNCTION", Token(MIN(), Nil) :: arg :: Nil) => Min(nodeToExpr(arg))
-    
+
     /* System functions about string operations */
     case Token("TOK_FUNCTION", Token(UPPER(), Nil) :: arg :: Nil) => Upper(nodeToExpr(arg))
     case Token("TOK_FUNCTION", Token(LOWER(), Nil) :: arg :: Nil) => Lower(nodeToExpr(arg))
-    
+
     /* Casts */
     case Token("TOK_FUNCTION", Token("TOK_STRING", Nil) :: arg :: Nil) =>
       Cast(nodeToExpr(arg), StringType)

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

@@ -32,7 +32,7 @@ import org.apache.hadoop.hive.serde2.avro.AvroSerDe
 
 import org.apache.spark.{SparkConf, SparkContext}
 import org.apache.spark.sql.catalyst.analysis._
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, NativeCommand}
+import org.apache.spark.sql.catalyst.plans.logical.{CacheCommand, LogicalPlan, NativeCommand}
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.hive._
 
@@ -103,7 +103,7 @@ class TestHiveContext(sc: SparkContext) extends LocalHiveContext(sc) {
   val inRepoTests = if (System.getProperty("user.dir").endsWith("sql" + File.separator + "hive")) {
     new File("src" + File.separator + "test" + File.separator + "resources" + File.separator)
   } else {
-    new File("sql" + File.separator + "hive" + File.separator + "src" + File.separator + "test" + 
+    new File("sql" + File.separator + "hive" + File.separator + "src" + File.separator + "test" +
       File.separator + "resources")
   }
 
@@ -130,6 +130,7 @@ class TestHiveContext(sc: SparkContext) extends LocalHiveContext(sc) {
     override lazy val analyzed = {
       val describedTables = logical match {
         case NativeCommand(describedTable(tbl)) => tbl :: Nil
+        case CacheCommand(tbl, _) => tbl :: Nil
         case _ => Nil
       }
 

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

@@ -56,4 +56,18 @@ class CachedTableSuite extends HiveComparisonTest {
       TestHive.uncacheTable("src")
     }
   }
+
+  test("'CACHE TABLE' and 'UNCACHE TABLE' HiveQL statement") {
+    TestHive.hql("CACHE TABLE src")
+    TestHive.table("src").queryExecution.executedPlan match {
+      case _: InMemoryColumnarTableScan => // Found evidence of caching
+      case _ => fail(s"Table 'src' should be cached")
+    }
+
+    TestHive.hql("UNCACHE TABLE src")
+    TestHive.table("src").queryExecution.executedPlan match {
+      case _: InMemoryColumnarTableScan => fail(s"Table 'src' should not be cached")
+      case _ => // Found evidence of uncaching
+    }
+  }
 }