[SPARK-2594][SQL] Add CACHE TABLE <name> AS SELECT ... (Updated as per review comments)

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

@@ -127,6 +127,7 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
   protected val SUBSTRING = Keyword("SUBSTRING")
   protected val SQRT = Keyword("SQRT")
   protected val ABS = Keyword("ABS")
+  protected val ADD = Keyword("ADD")
 
   // Use reflection to find the reserved words defined in this class.
   protected val reservedWords =
@@ -151,7 +152,7 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
         EXCEPT ^^^ { (q1: LogicalPlan, q2: LogicalPlan) => Except(q1, q2)} |
         UNION ~ opt(DISTINCT) ^^^ { (q1: LogicalPlan, q2: LogicalPlan) => Distinct(Union(q1, q2)) }
       )
-    | insert | cache
+    | insert | cache | unCache
   )
 
   protected lazy val select: Parser[LogicalPlan] =
@@ -181,11 +182,25 @@ class SqlParser extends StandardTokenParsers with PackratParsers {
         val overwrite: Boolean = o.getOrElse("") == "OVERWRITE"
         InsertIntoTable(r, Map[String, Option[String]](), s, overwrite)
     }
+
+  protected lazy val addCache: Parser[LogicalPlan] =
+    ADD ~ CACHE ~ TABLE ~> ident ~ AS ~ select <~ opt(";") ^^ {
+     case tableName ~ as ~ s =>
+       CacheTableAsSelectCommand(tableName,s)
+    }        
 
   protected lazy val cache: Parser[LogicalPlan] =
-    (CACHE ^^^ true | UNCACHE ^^^ false) ~ TABLE ~ ident ^^ {
-      case doCache ~ _ ~ tableName => CacheCommand(tableName, doCache)
+    CACHE ~ TABLE ~> ident ~ opt(AS) ~ opt(select) <~ opt(";") ^^ {
+      case tableName ~ None ~ None => 
+        CacheCommand(tableName, true)
+      case tableName ~ as ~ Some(plan) => 
+        CacheTableAsSelectCommand(tableName,plan)
     }
+
+  protected lazy val unCache: Parser[LogicalPlan] =
+    UNCACHE ~ TABLE ~> ident <~ opt(";") ^^ {
+      case tableName => CacheCommand(tableName, false)
+    }    
 
   protected lazy val projections: Parser[Seq[Expression]] = repsep(projection, ",")
 

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

@@ -75,3 +75,8 @@ case class DescribeCommand(
     AttributeReference("data_type", StringType, nullable = false)(),
     AttributeReference("comment", StringType, nullable = false)())
 }
+
+/**
+ * Returned for the "ADD CACHE TABLE tableName AS SELECT .." command.
+ */
+case class CacheTableAsSelectCommand(tableName: String, plan: LogicalPlan) extends Command

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

@@ -305,6 +305,8 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
         Seq(execution.ExplainCommand(logicalPlan, plan.output, extended)(context))
       case logical.CacheCommand(tableName, cache) =>
         Seq(execution.CacheCommand(tableName, cache)(context))
+      case logical.CacheTableAsSelectCommand(tableName,plan) =>
+        Seq(execution.CacheTableAsSelectCommand(tableName, plan))
       case _ => Nil
     }
   }

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

@@ -166,3 +166,22 @@ case class DescribeCommand(child: SparkPlan, output: Seq[Attribute])(
       child.output.map(field => Row(field.name, field.dataType.toString, null))
   }
 }
+
+/**
+ * :: DeveloperApi ::
+ */
+@DeveloperApi
+case class CacheTableAsSelectCommand(tableName: String, plan: LogicalPlan)
+  extends LeafNode with Command {
+
+  override protected[sql] lazy val sideEffectResult = {
+    sqlContext.catalog.registerTable(None, tableName,  sqlContext.executePlan(plan).analyzed)
+    sqlContext.cacheTable(tableName)
+    // It does the caching eager.
+    sqlContext.table(tableName).count
+    Seq.empty[Row]
+  }
+
+  override def output: Seq[Attribute] = Seq.empty  
+
+}

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

@@ -119,4 +119,16 @@ class CachedTableSuite extends QueryTest {
     }
     assert(!TestSQLContext.isCached("testData"), "Table 'testData' should not be cached")
   }
+
+  test("CACHE TABLE tableName AS SELECT Star Table") {
+    TestSQLContext.sql("CACHE TABLE testCacheTable AS SELECT * FROM testData")
+    TestSQLContext.sql("SELECT * FROM testCacheTable WHERE key = 1").collect()
+    TestSQLContext.uncacheTable("testCacheTable")
+  }
+
+  test("'CACHE TABLE tableName AS SELECT ..'") {
+    TestSQLContext.sql("CACHE TABLE testCacheTable AS SELECT * FROM testData")
+    assert(TestSQLContext.isCached("testCacheTable"), "Table 'testCacheTable' should be cached")
+    TestSQLContext.uncacheTable("testCacheTable")
+  }
 }

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

@@ -214,6 +214,7 @@ private[hive] object HiveQl {
    */
   def getAst(sql: String): ASTNode = ParseUtils.findRootNonNullToken((new ParseDriver).parse(sql))
 
+
   /** Returns a LogicalPlan for a given HiveQL string. */
   def parseSql(sql: String): LogicalPlan = {
     try {
@@ -229,11 +230,17 @@ private[hive] object HiveQl {
             SetCommand(Some(key), Some(value))
         }
       } else if (sql.trim.toLowerCase.startsWith("cache table")) {
-        CacheCommand(sql.trim.drop(12).trim, true)
+        sql.trim.drop(12).trim.split(" ").toSeq match {
+          case Seq(tableName) => 
+            CacheCommand(tableName, true)
+          case Seq(tableName,as, select@_*) => 
+            CacheTableAsSelectCommand(tableName,
+                createPlan(sql.trim.drop(12 + tableName.length() + as.length() + 2)))
+        }
       } else if (sql.trim.toLowerCase.startsWith("uncache table")) {
         CacheCommand(sql.trim.drop(14).trim, false)
       } else if (sql.trim.toLowerCase.startsWith("add jar")) {
-        AddJar(sql.trim.drop(8).trim)
+        NativeCommand(sql)
       } else if (sql.trim.toLowerCase.startsWith("add file")) {
         AddFile(sql.trim.drop(9))
       } else if (sql.trim.toLowerCase.startsWith("dfs")) {
@@ -243,15 +250,7 @@ private[hive] object HiveQl {
       } else if (sql.trim.startsWith("!")) {
         ShellCommand(sql.drop(1))
       } else {
-        val tree = getAst(sql)
-        if (nativeCommands contains tree.getText) {
-          NativeCommand(sql)
-        } else {
-          nodeToPlan(tree) match {
-            case NativePlaceholder => NativeCommand(sql)
-            case other => other
-          }
-        }
+        createPlan(sql)
       }
     } catch {
       case e: Exception => throw new ParseException(sql, e)
@@ -262,6 +261,19 @@ private[hive] object HiveQl {
         """.stripMargin)
     }
   }
+
+  /** Creates LogicalPlan for a given HiveQL string. */
+  def createPlan(sql: String) = {
+    val tree = getAst(sql)
+    if (nativeCommands contains tree.getText) {
+      NativeCommand(sql)
+    } else {
+      nodeToPlan(tree) match {
+        case NativePlaceholder => NativeCommand(sql)
+        case other => other
+      }
+    }
+  }
 
   def parseDdl(ddl: String): Seq[Attribute] = {
     val tree =
@@ -1097,7 +1109,7 @@ private[hive] object HiveQl {
 
       case Token("TOK_FUNCTION", Token(functionName, Nil) :: children) =>
         HiveGenericUdtf(functionName, attributes, children.map(nodeToExpr))
-
+        
       case a: ASTNode =>
         throw new NotImplementedError(
           s"""No parse rules for ASTNode type: ${a.getType}, text: ${a.getText}, tree: