[SPARK-33448][SQL] Support CACHE/UNCACHE TABLE commands for v2 tables

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -257,7 +257,7 @@ class Analyzer(override val catalogManager: CatalogManager)
       TypeCoercion.typeCoercionRules ++
       extendedResolutionRules : _*),
     Batch("Post-Hoc Resolution", Once,
-      Seq(ResolveNoopDropTable) ++
+      Seq(ResolveCommandsWithIfExists) ++
       postHocResolutionRules: _*),
     Batch("Normalize Alter Table", Once, ResolveAlterTableChanges),
     Batch("Remove Unresolved Hints", Once,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveNoopDropTable.scala → sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveCommandsWithIfExists.scala

@@ -17,17 +17,20 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.plans.logical.{DropTable, LogicalPlan, NoopDropTable}
+import org.apache.spark.sql.catalyst.plans.logical.{DropTable, LogicalPlan, NoopCommand, UncacheTable}
 import org.apache.spark.sql.catalyst.rules.Rule
 
 /**
- * A rule for handling [[DropTable]] logical plan when the table or temp view is not resolved.
- * If "ifExists" flag is set to true, the plan is resolved to [[NoopDropTable]],
+ * A rule for handling commands when the table or temp view is not resolved.
+ * These commands support a flag, "ifExists", so that they do not fail when a relation is not
+ * resolved. If the "ifExists" flag is set to true. the plan is resolved to [[NoopCommand]],
  * which is a no-op command.
  */
-object ResolveNoopDropTable extends Rule[LogicalPlan] {
+object ResolveCommandsWithIfExists extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsUp {
     case DropTable(u: UnresolvedTableOrView, ifExists, _) if ifExists =>
-      NoopDropTable(u.multipartIdentifier)
+      NoopCommand("DROP TABLE", u.multipartIdentifier)
+    case UncacheTable(u: UnresolvedTableOrView, ifExists) if ifExists =>
+      NoopCommand("UNCACHE TABLE", u.multipartIdentifier)
   }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -3323,7 +3323,7 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
   }
 
   /**
-   * Create a [[CacheTableStatement]].
+   * Create a [[CacheTable]] pr [[CacheTableAsSelect]].
    *
    * For example:
    * {{{
@@ -3343,14 +3343,21 @@ class AstBuilder extends SqlBaseBaseVisitor[AnyRef] with SQLConfHelper with Logg
         "the table name in CACHE TABLE AS SELECT", ctx)
     }
     val options = Option(ctx.options).map(visitPropertyKeyValues).getOrElse(Map.empty)
-    CacheTableStatement(tableName, query, ctx.LAZY != null, options)
+    val isLazy = ctx.LAZY != null
+    if (query.isDefined) {
+      CacheTableAsSelect(tableName.head, query.get, isLazy, options)
+    } else {
+      CacheTable(UnresolvedTableOrView(tableName), isLazy, options)
+    }
   }
 
   /**
-   * Create an [[UncacheTableStatement]] logical plan.
+   * Create an [[UncacheTable]] logical plan.
    */
   override def visitUncacheTable(ctx: UncacheTableContext): LogicalPlan = withOrigin(ctx) {
-    UncacheTableStatement(visitMultipartIdentifier(ctx.multipartIdentifier), ctx.EXISTS != null)
+    UncacheTable(
+      UnresolvedTableOrView(visitMultipartIdentifier(ctx.multipartIdentifier)),
+      ctx.EXISTS != null)
   }
 
   /**

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

@@ -329,22 +329,6 @@ case class UseStatement(isNamespaceSet: Boolean, nameParts: Seq[String]) extends
  */
 case class RepairTableStatement(tableName: Seq[String]) extends ParsedStatement
 
-/**
- * A CACHE TABLE statement, as parsed from SQL
- */
-case class CacheTableStatement(
-    tableName: Seq[String],
-    plan: Option[LogicalPlan],
-    isLazy: Boolean,
-    options: Map[String, String]) extends ParsedStatement
-
-/**
- * An UNCACHE TABLE statement, as parsed from SQL
- */
-case class UncacheTableStatement(
-    tableName: Seq[String],
-    ifExists: Boolean) extends ParsedStatement
-
 /**
  * A TRUNCATE TABLE statement, as parsed from SQL
  */

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

@@ -417,9 +417,9 @@ case class DropTable(
 }
 
 /**
- * The logical plan for handling non-existing table for DROP TABLE command.
+ * The logical plan for no-op command handling non-existing table.
  */
-case class NoopDropTable(multipartIdentifier: Seq[String]) extends Command
+case class NoopCommand(commandName: String, multipartIdentifier: Seq[String]) extends Command
 
 /**
  * The logical plan of the ALTER TABLE command.
@@ -670,3 +670,29 @@ case class LoadData(
 case class ShowCreateTable(child: LogicalPlan, asSerde: Boolean = false) extends Command {
   override def children: Seq[LogicalPlan] = child :: Nil
 }
+
+/**
+ * The logical plan of the CACHE TABLE command.
+ */
+case class CacheTable(
+    child: LogicalPlan,
+    isLazy: Boolean,
+    options: Map[String, String]) extends Command {
+  override def children: Seq[LogicalPlan] = child :: Nil
+}
+
+/**
+ * The logical plan of the CACHE TABLE ... AS SELECT command.
+ */
+case class CacheTableAsSelect(
+    tempViewName: String,
+    plan: LogicalPlan,
+    isLazy: Boolean,
+    options: Map[String, String]) extends Command
+
+/**
+ * The logical plan of the UNCACHE TABLE command.
+ */
+case class UncacheTable(child: LogicalPlan, ifExists: Boolean) extends Command {
+  override def children: Seq[LogicalPlan] = child :: Nil
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/DDLParserSuite.scala

@@ -1594,15 +1594,26 @@ class DDLParserSuite extends AnalysisTest {
   test("CACHE TABLE") {
     comparePlans(
       parsePlan("CACHE TABLE a.b.c"),
-      CacheTableStatement(Seq("a", "b", "c"), None, false, Map.empty))
+      CacheTable(UnresolvedTableOrView(Seq("a", "b", "c")), false, Map.empty))
+
+    comparePlans(
+      parsePlan("CACHE TABLE t AS SELECT * FROM testData"),
+      CacheTableAsSelect(
+        "t",
+        Project(Seq(UnresolvedStar(None)), UnresolvedRelation(Seq("testData"))),
+        false,
+        Map.empty))
 
     comparePlans(
       parsePlan("CACHE LAZY TABLE a.b.c"),
-      CacheTableStatement(Seq("a", "b", "c"), None, true, Map.empty))
+      CacheTable(UnresolvedTableOrView(Seq("a", "b", "c")), true, Map.empty))
 
     comparePlans(
       parsePlan("CACHE LAZY TABLE a.b.c OPTIONS('storageLevel' 'DISK_ONLY')"),
-      CacheTableStatement(Seq("a", "b", "c"), None, true, Map("storageLevel" -> "DISK_ONLY")))
+      CacheTable(
+        UnresolvedTableOrView(Seq("a", "b", "c")),
+        true,
+        Map("storageLevel" -> "DISK_ONLY")))
 
     intercept("CACHE TABLE a.b.c AS SELECT * FROM testData",
       "It is not allowed to add catalog/namespace prefix a.b")
@@ -1611,11 +1622,11 @@ class DDLParserSuite extends AnalysisTest {
   test("UNCACHE TABLE") {
     comparePlans(
       parsePlan("UNCACHE TABLE a.b.c"),
-      UncacheTableStatement(Seq("a", "b", "c"), ifExists = false))
+      UncacheTable(UnresolvedTableOrView(Seq("a", "b", "c")), ifExists = false))
 
     comparePlans(
       parsePlan("UNCACHE TABLE IF EXISTS a.b.c"),
-      UncacheTableStatement(Seq("a", "b", "c"), ifExists = true))
+      UncacheTable(UnresolvedTableOrView(Seq("a", "b", "c")), ifExists = true))
   }
 
   test("TRUNCATE table") {

sql/core/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveSessionCatalog.scala

@@ -442,19 +442,15 @@ class ResolveSessionCatalog(
         ShowCreateTableCommand(ident.asTableIdentifier)
       }
 
-    case CacheTableStatement(tbl, plan, isLazy, options) =>
-      val name = if (plan.isDefined) {
-        // CACHE TABLE ... AS SELECT creates a temp view with the input query.
-        // Temp view doesn't belong to any catalog and we shouldn't resolve catalog in the name.
-        tbl
-      } else {
-        parseTempViewOrV1Table(tbl, "CACHE TABLE")
-      }
-      CacheTableCommand(name.asTableIdentifier, plan, isLazy, options)
+    // CACHE TABLE ... AS SELECT creates a temp view with the input query.
+    case CacheTableAsSelect(tempViewName, plan, isLazy, options) =>
+      CacheTableCommand(TableIdentifier(tempViewName), Some(plan), isLazy, options)
 
-    case UncacheTableStatement(tbl, ifExists) =>
-      val name = parseTempViewOrV1Table(tbl, "UNCACHE TABLE")
-      UncacheTableCommand(name.asTableIdentifier, ifExists)
+    case CacheTable(ResolvedV1TableOrViewIdentifier(ident), isLazy, options) =>
+      CacheTableCommand(ident.asTableIdentifier, None, isLazy, options)
+
+    case UncacheTable(ResolvedV1TableOrViewIdentifier(ident), ifExists) =>
+      UncacheTableCommand(ident.asTableIdentifier, ifExists)
 
     case TruncateTableStatement(tbl, partitionSpec) =>
       val v1TableName = parseV1Table(tbl, "TRUNCATE TABLE")
@@ -570,12 +566,9 @@ class ResolveSessionCatalog(
             "SHOW VIEWS, only SessionCatalog supports this command.")
       }
 
-    case ShowTableProperties(ResolvedV1TableIdentifier(ident), propertyKey) =>
+    case ShowTableProperties(ResolvedV1TableOrViewIdentifier(ident), propertyKey) =>
       ShowTablePropertiesCommand(ident.asTableIdentifier, propertyKey)
 
-    case ShowTableProperties(r: ResolvedView, propertyKey) =>
-      ShowTablePropertiesCommand(r.identifier.asTableIdentifier, propertyKey)
-
     case DescribeFunction(ResolvedFunc(identifier), extended) =>
       DescribeFunctionCommand(identifier.asFunctionIdentifier, extended)
 

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

@@ -0,0 +1,36 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution
+
+import java.util.Locale
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
+
+object CacheTableUtils extends Logging {
+  def getStorageLevel(options: Map[String, String]): Option[String] = {
+    val storageLevelKey = "storagelevel"
+    val storageLevelValue =
+      CaseInsensitiveMap(options).get(storageLevelKey).map(_.toUpperCase(Locale.ROOT))
+    val withoutStorageLevel = options.filterKeys(_.toLowerCase(Locale.ROOT) != storageLevelKey)
+    if (withoutStorageLevel.nonEmpty) {
+      logWarning(s"Invalid options: ${withoutStorageLevel.mkString(", ")}")
+    }
+    storageLevelValue
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/command/cache.scala

@@ -17,13 +17,11 @@
 
 package org.apache.spark.sql.execution.command
 
-import java.util.Locale
-
 import org.apache.spark.sql.{Dataset, Row, SparkSession}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.{IgnoreCachedData, LogicalPlan}
-import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
+import org.apache.spark.sql.execution.CacheTableUtils
 import org.apache.spark.storage.StorageLevel
 
 case class CacheTableCommand(
@@ -41,17 +39,10 @@ case class CacheTableCommand(
       Dataset.ofRows(sparkSession, logicalPlan).createTempView(tableIdent.quotedString)
     }
 
-    val storageLevelKey = "storagelevel"
-    val storageLevelValue =
-      CaseInsensitiveMap(options).get(storageLevelKey).map(_.toUpperCase(Locale.ROOT))
-    val withoutStorageLevel = options.filterKeys(_.toLowerCase(Locale.ROOT) != storageLevelKey)
-    if (withoutStorageLevel.nonEmpty) {
-      logWarning(s"Invalid options: ${withoutStorageLevel.mkString(", ")}")
-    }
-
-    if (storageLevelValue.nonEmpty) {
+    val optStorageLevel = CacheTableUtils.getStorageLevel(options)
+    if (optStorageLevel.nonEmpty) {
       sparkSession.catalog.cacheTable(
-        tableIdent.quotedString, StorageLevel.fromString(storageLevelValue.get))
+        tableIdent.quotedString, StorageLevel.fromString(optStorageLevel.get))
     } else {
       sparkSession.catalog.cacheTable(tableIdent.quotedString)
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/CacheTableExec.scala

@@ -0,0 +1,79 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2
+
+import org.apache.spark.sql.{Dataset, SparkSession}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.connector.catalog.{Identifier, Table, TableCatalog}
+import org.apache.spark.sql.execution.CacheTableUtils
+import org.apache.spark.storage.StorageLevel
+
+/**
+ * Physical plan node for caching a table.
+ */
+case class CacheTableExec(
+    session: SparkSession,
+    catalog: TableCatalog,
+    table: Table,
+    ident: Identifier,
+    isLazy: Boolean,
+    options: Map[String, String]) extends V2CommandExec {
+  override def run(): Seq[InternalRow] = {
+    import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.IdentifierHelper
+
+    val v2Relation = DataSourceV2Relation.create(table, Some(catalog), Some(ident))
+    val df = Dataset.ofRows(session, v2Relation)
+    val tableName = Some(ident.quoted)
+    val optStorageLevel = CacheTableUtils.getStorageLevel(options)
+    if (optStorageLevel.nonEmpty) {
+      session.sharedState.cacheManager.cacheQuery(
+        df, tableName, StorageLevel.fromString(optStorageLevel.get))
+    } else {
+      session.sharedState.cacheManager.cacheQuery(df, tableName)
+    }
+
+    if (!isLazy) {
+      // Performs eager caching.
+      df.count()
+    }
+
+    Seq.empty
+  }
+
+  override def output: Seq[Attribute] = Seq.empty
+}
+
+/**
+ * Physical plan node for uncaching a table.
+ */
+case class UncacheTableExec(
+    session: SparkSession,
+    catalog: TableCatalog,
+    table: Table,
+    ident: Identifier) extends V2CommandExec {
+  override def run(): Seq[InternalRow] = {
+    val v2Relation = DataSourceV2Relation.create(table, Some(catalog), Some(ident))
+    val df = Dataset.ofRows(session, v2Relation)
+    // Cascade should be true unless a temporary view is uncached.
+    session.sharedState.cacheManager.uncacheQuery(df, cascade = true)
+    Seq.empty
+  }
+
+  override def output: Seq[Attribute] = Seq.empty
+}

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala

@@ -234,7 +234,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
     case DropTable(r: ResolvedTable, ifExists, purge) =>
       DropTableExec(session, r.catalog, r.table, r.identifier, ifExists, purge) :: Nil
 
-    case _: NoopDropTable =>
+    case _: NoopCommand =>
       LocalTableScanExec(Nil, Nil) :: Nil
 
     case AlterTable(catalog, ident, _, changes) =>
@@ -302,6 +302,12 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat
     case ShowCreateTable(_: ResolvedTable, _) =>
       throw new AnalysisException("SHOW CREATE TABLE is not supported for v2 tables.")
 
+    case CacheTable(r: ResolvedTable, isLazy, options) =>
+      CacheTableExec(session, r.catalog, r.table, r.identifier, isLazy, options) :: Nil
+
+    case UncacheTable(r: ResolvedTable, _) =>
+      UncacheTableExec(session, r.catalog, r.table, r.identifier) :: Nil
+
     case _ => Nil
   }
 }