diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsTableStateOptions.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsTableStateOptions.java new file mode 100644 index 0000000000000..d30c37ec81a2b --- /dev/null +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsTableStateOptions.java @@ -0,0 +1,52 @@ +/* + * 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.connector.catalog; + +import java.util.Set; + +import org.apache.spark.annotation.Evolving; + +/** + * A catalog capability for identifying options that select a table's state. + *

+ * Spark may resolve the same table more than once while analyzing or refreshing one query. A + * catalog can implement this interface to declare which raw read options may cause + * {@link TableCatalog#loadTable(Identifier, TableContext, + * org.apache.spark.sql.util.CaseInsensitiveStringMap)} to select a different table state, such as + * a branch, tag, snapshot, or version. Spark can then reuse one concrete {@link Table} instance + * for references whose table-state options match while preserving every reference's complete + * option map for scan planning. + *

+ * Option key matching is case-insensitive. Option values remain case-sensitive. Parsed Spark time + * travel is handled independently and must not be included in the returned set. + *

+ * Catalogs that do not implement this capability are handled conservatively: Spark treats every + * raw option as table-state-affecting. + * + * @since 4.3.0 + */ +@Evolving +public interface SupportsTableStateOptions extends CatalogPlugin { + + /** + * Returns the raw option keys that may affect the table state selected by {@code loadTable}. + * + * @return a non-null set of case-insensitive option keys + */ + Set tableStateOptionKeys(); +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala index 4b2a1cc1114c9..e790ea2e61a88 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala @@ -135,6 +135,9 @@ object FakeV2SessionCatalog extends TableCatalog with FunctionCatalog with Suppo * @param relationCache A mapping from (qualified table name, time travel spec, options) to * resolved relations. This can ensure that the table is resolved only once if * a table is used multiple times in a query with the same options. + * @param tableCache A mapping from (catalog, identifier, time travel spec, table-state options) to + * concrete tables. This pins one table state while allowing references to keep + * different read-specific options. * @param referredTempViewNames All the temp view names referred by the current view we are * resolving. It's used to make sure the relation resolution is * consistent between view creation and view resolution. For example, @@ -155,6 +158,7 @@ case class AnalysisContext( nestedViewDepth: Int = 0, maxNestedViewDepth: Int = -1, relationCache: mutable.Map[RelationCacheKey, LogicalPlan] = mutable.Map.empty, + tableCache: mutable.Map[TableCacheKey, Table] = mutable.Map.empty, referredTempViewNames: Seq[Seq[String]] = Seq.empty, // 1. If we are resolving a view, this field will be restored from the view metadata, // by calling `AnalysisContext.withAnalysisContext(viewDesc)`. @@ -249,6 +253,7 @@ object AnalysisContext { nestedViewDepth = originContext.nestedViewDepth + 1, maxNestedViewDepth = maxNestedViewDepth, relationCache = originContext.relationCache, + tableCache = originContext.tableCache, referredTempViewNames = viewDesc.viewReferredTempViewNames, referredTempFunctionNames = mutable.Set(viewDesc.viewReferredTempFunctionNames: _*), referredTempVariableNames = viewDesc.viewReferredTempVariableNames, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala index d217219cd8f95..424e66c8225ab 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RelationResolution.scala @@ -62,6 +62,7 @@ class RelationResolution( val v1SessionCatalog = catalogManager.v1SessionCatalog private def relationCache = AnalysisContext.get.relationCache + private def tableCache = AnalysisContext.get.tableCache /** * If we are resolving database objects (relations, functions, etc.) inside views, we may need to @@ -254,6 +255,10 @@ class RelationResolution( cached .map(adaptCachedRelation(_, planId)) .orElse { + lazy val tableKey = + toTableCacheKey(catalog, ident, finalTimeTravelSpec, finalOptions) + val pinnedTable = if (writePrivileges == null) tableCache.get(tableKey) else None + // For a `RelationCatalog` with no time-travel / write privileges, the single-RPC // `loadRelation` answers both "is there a table?" and "is there a view?" in one // call. Time-travel and write privileges apply to tables only, so for those the @@ -263,45 +268,48 @@ class RelationResolution( // Skip the table-side lookup entirely for view-only catalogs (no `TableCatalog` // mixin): `CatalogV2Util.loadTable` would call `asTableCatalog` and throw // MISSING_CATALOG_ABILITY.TABLES, masking the legitimate view-resolution path. - val relation: Option[Relation] = catalog match { - case mc: RelationCatalog if finalTimeTravelSpec.isEmpty && writePrivileges == null => - try { - Some(mc.loadRelation(ident)) - } catch { - case _: NoSuchTableException => None - } - case _ => - val tableSide: Option[Table] = if ( - CatalogV2Util.isSessionCatalog(catalog) || catalog.isInstanceOf[TableCatalog] - ) { - CatalogV2Util.loadTable( - catalog, - ident, - finalTimeTravelSpec, - Option(writePrivileges), - finalOptions) - } else { - None - } - // Fallback to ViewCatalog for catalogs that host views but where loadTable - // returned None (or was skipped because there's no TableCatalog mixin). - // Time-travel / write privileges only apply to tables, not views, so the - // fallback only fires when both are absent. - tableSide.orElse { - if (finalTimeTravelSpec.isEmpty && writePrivileges == null) { - catalog match { - case vc: ViewCatalog => - try { - Some(vc.loadView(ident)) - } catch { - case _: NoSuchViewException => None - } - case _ => None - } + val relation: Option[Relation] = pinnedTable.orElse { + catalog match { + case mc: RelationCatalog + if finalTimeTravelSpec.isEmpty && writePrivileges == null => + try { + Some(mc.loadRelation(ident)) + } catch { + case _: NoSuchTableException => None + } + case _ => + val tableSide: Option[Table] = if ( + CatalogV2Util.isSessionCatalog(catalog) || catalog.isInstanceOf[TableCatalog] + ) { + CatalogV2Util.loadTable( + catalog, + ident, + finalTimeTravelSpec, + Option(writePrivileges), + finalOptions) } else { None } - } + // Fallback to ViewCatalog for catalogs that host views but where loadTable + // returned None (or was skipped because there's no TableCatalog mixin). + // Time-travel / write privileges only apply to tables, not views, so the + // fallback only fires when both are absent. + tableSide.orElse { + if (finalTimeTravelSpec.isEmpty && writePrivileges == null) { + catalog match { + case vc: ViewCatalog => + try { + Some(vc.loadView(ident)) + } catch { + case _: NoSuchViewException => None + } + case _ => None + } + } else { + None + } + } + } } // `table` is `relation` filtered to tables only -- used for cache lookup since // we don't share-cache views. @@ -312,12 +320,14 @@ class RelationResolution( // `Table`. val sharedRelationCacheMatch = for { t <- table - if finalTimeTravelSpec.isEmpty && writePrivileges == null && !u.isStreaming + if pinnedTable.isEmpty && finalTimeTravelSpec.isEmpty && + writePrivileges == null && !u.isStreaming cached <- lookupSharedRelationCache(catalog, ident, t) if cached.options == finalOptions } yield { val nameParts = ident.toQualifiedNameParts(catalog) val aliasedRelation = SubqueryAlias(nameParts, cached) + tableCache.update(tableKey, cached.table) relationCache.update(key, aliasedRelation) adaptCachedRelation(aliasedRelation, planId) } @@ -330,6 +340,9 @@ class RelationResolution( finalOptions, u.isStreaming, finalTimeTravelSpec) + if (writePrivileges == null && pinnedTable.isEmpty) { + table.foreach(tableCache.update(tableKey, _)) + } loaded.foreach(relationCache.update(key, _)) loaded.map(cloneWithPlanId(_, planId)) } @@ -476,7 +489,12 @@ class RelationResolution( def resolveReference(ref: V2TableReference): LogicalPlan = { val relation = if (ref.context.cacheable) { - getOrLoadRelation(ref) + // A temporary view may contain a relation pinned by CacheManager, so its re-resolution + // consults sharedRelationCache to preserve that Table. Transaction references use the + // Table loaded through the transaction catalog instead. + val useSharedRelationCache = + ref.context.isInstanceOf[V2TableReference.TemporaryViewContext] + getOrLoadRelation(ref, useSharedRelationCache) } else { loadRelation(ref) } @@ -484,15 +502,47 @@ class RelationResolution( cloneWithPlanId(relation, planId) } - private def getOrLoadRelation(ref: V2TableReference): LogicalPlan = { + private def getOrLoadRelation( + ref: V2TableReference, + useSharedRelationCache: Boolean): LogicalPlan = { val key = toCacheKey(ref.catalog, ref.identifier, None, ref.options) relationCache.get(key) match { case Some(cached) => adaptCachedRelation(cached, ref) case None => - val relation = loadRelation(ref) - relationCache.update(key, relation) - relation + val resolvedCatalog = catalogManager.catalog(ref.catalog.name).asTableCatalog + val tableKey = toTableCacheKey(resolvedCatalog, ref.identifier, None, ref.options) + tableCache.get(tableKey) match { + case Some(pinnedTable) => + val relation = createRelation(ref, resolvedCatalog, pinnedTable) + relationCache.update(key, relation) + relation + case None => + val loadedTable = CatalogV2Util.getTable( + resolvedCatalog, + ref.identifier, + options = ref.options) + val sharedRelationCacheMatch = if (useSharedRelationCache) { + lookupSharedRelationCache( + resolvedCatalog, + ref.identifier, + loadedTable).filter(_.options == ref.options) + } else { + None + } + sharedRelationCacheMatch match { + case Some(cached) => + val relation = adaptCachedRelation(cached, ref) + tableCache.update(tableKey, cached.table) + relationCache.update(key, relation) + relation + case None => + val relation = createRelation(ref, resolvedCatalog, loadedTable) + tableCache.update(tableKey, loadedTable) + relationCache.update(key, relation) + relation + } + } } } @@ -508,6 +558,13 @@ class RelationResolution( private def loadRelation(ref: V2TableReference): LogicalPlan = { val resolvedCatalog = catalogManager.catalog(ref.catalog.name).asTableCatalog val table = resolvedCatalog.loadTable(ref.identifier) + createRelation(ref, resolvedCatalog, table) + } + + private def createRelation( + ref: V2TableReference, + resolvedCatalog: TableCatalog, + table: Table): DataSourceV2Relation = { V2TableReferenceUtils.validateLoadedTable(table, ref) DataSourceV2Relation( table = table, @@ -551,6 +608,18 @@ class RelationResolution( RelationCacheKey(nameParts, timeTravelSpec, options) } + private def toTableCacheKey( + catalog: CatalogPlugin, + ident: Identifier, + timeTravelSpec: Option[TimeTravelSpec], + options: CaseInsensitiveStringMap): TableCacheKey = { + TableCacheKey( + catalog, + ident, + timeTravelSpec, + CatalogV2Util.tableStateOptions(catalog, options)) + } + private def cloneWithPlanId(plan: LogicalPlan, planId: Option[Long]): LogicalPlan = { planId match { case Some(id) => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableCacheKey.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableCacheKey.scala new file mode 100644 index 0000000000000..186c4b9855a43 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/TableCacheKey.scala @@ -0,0 +1,33 @@ +/* + * 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.catalyst.analysis + +import org.apache.spark.sql.connector.catalog.{CatalogPlugin, Identifier} +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Key for the per-query table-state cache in [[AnalysisContext]]. + * + * Unlike [[RelationCacheKey]], this key contains only options declared to affect table state. This + * lets references retain different scan options while sharing one concrete table state. + */ +private[sql] case class TableCacheKey( + catalog: CatalogPlugin, + identifier: Identifier, + timeTravelSpec: Option[TimeTravelSpec], + stateOptions: CaseInsensitiveStringMap) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala index 28aaae4a81ffa..b57536c8f20c5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogV2Util.scala @@ -481,6 +481,29 @@ private[sql] object CatalogV2Util { case _: NoSuchDatabaseException => None } + /** + * Projects a complete read option map to the options that may select table state. + * + * Catalogs must explicitly opt in to projection. For all other catalogs, every option is kept + * so that reusing a concrete table cannot silently combine states the catalog considers + * different. + */ + def tableStateOptions( + catalog: CatalogPlugin, + options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = catalog match { + case supports: SupportsTableStateOptions => + val stateKeys = supports.tableStateOptionKeys().asScala + .map(_.toLowerCase(Locale.ROOT)) + .toSet + val projected = options.entrySet().asScala.collect { + case entry if stateKeys.contains(entry.getKey.toLowerCase(Locale.ROOT)) => + entry.getKey -> entry.getValue + }.toMap + new CaseInsensitiveStringMap(projected.asJava) + case _ => + options + } + def getTable( catalog: CatalogPlugin, ident: Identifier, diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala index 75846aa49616c..c078504ea33a6 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/TableLookupCacheSuite.scala @@ -119,4 +119,21 @@ class TableLookupCacheSuite extends AnalysisTest with Matchers { verify(catalog, times(1)).getTable("default", "t1") } } + + test("nested view analysis shares both query-scoped caches") { + AnalysisContext.withNewAnalysisContext { + val outer = AnalysisContext.get + val viewDesc = CatalogTable( + TableIdentifier("view", Some("default")), + CatalogTableType.VIEW, + CatalogStorageFormat.empty, + StructType(Seq(StructField("a", IntegerType))), + viewText = Some("select * from t1")) + + AnalysisContext.withAnalysisContext(viewDesc) { + assert(AnalysisContext.get.relationCache eq outer.relationCache) + assert(AnalysisContext.get.tableCache eq outer.tableCache) + } + } + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala index 4b9d55be07e3f..c4c7b339ed9e7 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/CatalogV2UtilSuite.scala @@ -28,6 +28,14 @@ import org.apache.spark.sql.util.CaseInsensitiveStringMap class CatalogV2UtilSuite extends SparkFunSuite { + private def catalogWithStateOptions(keys: java.util.Set[String]): SupportsTableStateOptions = { + new SupportsTableStateOptions { + override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = {} + override def name(): String = "state-options" + override def tableStateOptionKeys(): java.util.Set[String] = keys + } + } + // CatalogV2Util.getTable routes through the options-aware TableCatalog.loadTable, whose default // implementation dispatches to the existing overloads. Stub only that method to run the real // default so the dispatch is exercised; the leaf overloads stay as plain mock methods (returning @@ -113,4 +121,29 @@ class CatalogV2UtilSuite extends SparkFunSuite { assert(a.toString.contains("timeTravel")) assert(a.toString.contains("writePrivileges")) } + + test("tableStateOptions projects declared keys case-insensitively") { + val catalog = catalogWithStateOptions(java.util.Set.of("BrAnCh", "tag")) + val options = new CaseInsensitiveStringMap(java.util.Map.of( + "branch", "Main", + "TAG", "Release", + "split-size", "5")) + + val stateOptions = CatalogV2Util.tableStateOptions(catalog, options) + + assert(stateOptions.size() == 2) + assert(stateOptions.get("BRANCH") == "Main") + assert(stateOptions.get("tag") == "Release") + assert(!stateOptions.containsKey("split-size")) + } + + test("tableStateOptions conservatively keeps all options for catalogs without capability") { + val catalog = mock(classOf[CatalogPlugin]) + val options = new CaseInsensitiveStringMap( + java.util.Map.of("branch", "Main", "split-size", "5")) + + val stateOptions = CatalogV2Util.tableStateOptions(catalog, options) + + assert(stateOptions eq options) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala index 10408aa513631..aa1f4c79aaedd 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2TableRefreshUtil.scala @@ -88,7 +88,8 @@ private[sql] object V2TableRefreshUtil extends SQLConfHelper with Logging { plan transformWithSubqueries { case r @ ExtractV2CatalogAndIdentifier(catalog, ident) if (r.isVersioned || !versionedOnly) && r.timeTravelSpec.isEmpty => - val currentTable = currentTables.getOrElseUpdate((catalog, ident, r.options), { + val stateOptions = CatalogV2Util.tableStateOptions(catalog, r.options) + val currentTable = currentTables.getOrElseUpdate((catalog, ident, stateOptions), { val tableName = V2TableUtil.toQualifiedName(catalog, ident) lookupCachedRelation(spark, catalog, ident, r.table) match { case Some(cached) if cached.options == r.options => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala index 8cdf12b18f449..b2296eba09759 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2OptionSuite.scala @@ -17,17 +17,32 @@ package org.apache.spark.sql.connector +import java.util import java.util.concurrent.atomic.AtomicInteger -import org.apache.spark.sql.{AnalysisException, Row} +import org.apache.spark.sql.{AnalysisException, DataFrame, Row} import org.apache.spark.sql.QueryTest.withQueryExecutionsCaptured -import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation +import org.apache.spark.sql.catalyst.analysis.{ + AnalysisContext, + RelationCache, + RelationResolution, + UnresolvedRelation, + V2TableReference} import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 -import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryBaseTable, InMemoryCatalog, InMemoryRowLevelOperationTableCatalog, Table, TimeTravel} +import org.apache.spark.sql.connector.catalog.{ + Identifier, + InMemoryBaseTable, + InMemoryCatalog, + InMemoryRowLevelOperationTableCatalog, + SupportsTableStateOptions, + Table, + TableWritePrivilege, + TimeTravel} import org.apache.spark.sql.execution.CommandResultExec import org.apache.spark.sql.execution.datasources.v2._ import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.util.CaseInsensitiveStringMap class LoadCountingInMemoryCatalog extends InMemoryCatalog { val singleArgLoads = new AtomicInteger(0) @@ -38,6 +53,11 @@ class LoadCountingInMemoryCatalog extends InMemoryCatalog { } } +class StateAwareInMemoryCatalog extends LoadCountingInMemoryCatalog + with SupportsTableStateOptions { + override def tableStateOptionKeys(): util.Set[String] = util.Set.of("snapshot") +} + class DataSourceV2OptionSuite extends DatasourceV2SQLBase { import testImplicits._ @@ -46,6 +66,20 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { private def inMemoryCatalog: InMemoryCatalog = catalog("testcat").asInstanceOf[InMemoryCatalog] + private def withStateAwareTable( + f: (StateAwareInMemoryCatalog, String) => Unit): Unit = { + withSQLConf( + "spark.sql.catalog.statecat" -> classOf[StateAwareInMemoryCatalog].getName, + "spark.sql.catalog.statecat.copyOnLoad" -> "true") { + val tableName = "statecat.ns.table" + withTable(tableName) { + sql(s"CREATE TABLE $tableName (id bigint, data string)") + sql(s"INSERT INTO $tableName VALUES (1, 'a'), (2, 'b')") + f(catalog("statecat").asInstanceOf[StateAwareInMemoryCatalog], tableName) + } + } + } + test("SPARK-36680: Supports Dynamic Table Options for SQL Select") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { @@ -548,17 +582,15 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { } } - test("SPARK-58389: a self-join with different options loads the table once per option bag") { + test("catalogs without state-option support load the table once per complete option bag") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { sql(s"CREATE TABLE $t1 (id bigint, data string)") sql(s"INSERT INTO $t1 VALUES (1, 'a'), (2, 'b')") inMemoryCatalog.resetLoadTableCalls() - // The two references share a name but carry different options. Because a catalog's - // options-aware loadTable can return a different Table depending on the options, the analyzer - // relation cache is keyed on the options, so each reference triggers its own loadTable rather - // than reusing the first reference's Table. + // The catalog does not classify its options, so Spark conservatively treats both complete + // option bags as different table states. val df = sql(s"SELECT a.id FROM $t1 WITH (`split-size` = 5) a " + s"JOIN $t1 WITH (`split-size` = 9) b ON a.id = b.id") df.queryExecution.analyzed @@ -590,6 +622,338 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { } } + test("same table state shares one Table while preserving each reference's options") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + + val df = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`SnApShOt` = 's1', `split-size` = 5) a JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON a.id = b.id") + + val analyzedRelations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("split-size") => r + } + assert(analyzedRelations.size == 2) + assert(analyzedRelations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(analyzedRelations.map(_.options.get("snapshot")).distinct == Seq("s1")) + assert(analyzedRelations.head.table eq analyzedRelations.last.table) + + val analysisLoads = stateCatalog.loadTableCalls.filter(_._2.get("snapshot") == "s1") + assert(analysisLoads.size == 1, + s"expected one catalog load for state s1 during analysis, got: $analysisLoads") + + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + assert(df.collect().toSeq == Seq(Row(1), Row(2))) + + val refreshLoads = stateCatalog.loadTableCalls.filter(_._2.get("snapshot") == "s1") + assert(refreshLoads.size == 1, + s"expected one catalog load for state s1 during refresh, got: $refreshLoads") + val refreshedRelations = df.queryExecution.optimizedPlan.collect { + case s: DataSourceV2ScanRelation if s.relation.options.containsKey("split-size") => + s.relation + } + assert(refreshedRelations.size == 2) + assert(refreshedRelations.head.table eq refreshedRelations.last.table) + } + } + + test("shared relation cache and refresh preserve first-resolution-wins table state") { + withStateAwareTable { (stateCatalog, tableName) => + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + cached.cache() + try { + cached.collect() + val cachedTable = cached.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r.table + }.getOrElse(fail("expected a cached v2 relation")) + + def relations(df: DataFrame): Seq[DataSourceV2Relation] = { + df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + } + + stateCatalog.resetLoadTableCalls() + val cachedFirst = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) a JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON a.id = b.id") + val cachedFirstRelations = relations(cachedFirst) + assert(cachedFirstRelations.size == 2) + assert(cachedFirstRelations.forall(_.table eq cachedTable)) + assert(stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") == 1) + + stateCatalog.resetLoadTableCalls() + assert(cachedFirst.collect().map(_.getLong(0)).sorted.toSeq == Seq(1L, 2L)) + assert(stateCatalog.loadTableCalls.isEmpty, + s"shared relation cache pin should avoid refresh reloads, got: " + + stateCatalog.loadTableCalls) + + stateCatalog.resetLoadTableCalls() + val uncachedFirst = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) a JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) b ON a.id = b.id") + val uncachedFirstRelations = relations(uncachedFirst) + assert(uncachedFirstRelations.size == 2) + assert(uncachedFirstRelations.head.table eq uncachedFirstRelations.last.table) + assert(uncachedFirstRelations.forall(_.table ne cachedTable)) + assert(stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") == 1) + + stateCatalog.resetLoadTableCalls() + assert(uncachedFirst.collect().map(_.getLong(0)).sorted.toSeq == Seq(1L, 2L)) + assert(stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") == 1) + } finally { + cached.unpersist() + } + } + } + + test("streaming references participate in the query table-state cache") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + val df = sql(s"SELECT a.id FROM STREAM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) a JOIN STREAM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON a.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: StreamingRelationV2 if r.extraOptions.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.extraOptions.get("split-size")).sorted == Seq("5", "9")) + assert(relations.head.table eq relations.last.table) + val stateLoads = stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") + assert(stateLoads == 1, s"expected one streaming table load, got: $stateLoads") + } + } + + test("different table-state option values establish separate table pins") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + + val df = sql(s"SELECT a.id FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5) a JOIN $tableName " + + s"WITH (`snapshot` = 's2', `split-size` = 9) b ON a.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.options.get("snapshot")).sorted == Seq("s1", "s2")) + assert(relations.head.table ne relations.last.table) + val loadedStates = stateCatalog.loadTableCalls + .map(_._2.get("snapshot")) + .filter(_ != null) + .sorted + assert(loadedStates == Seq("s1", "s2")) + } + } + + test("persistent write targets bypass both query-scoped read caches") { + withStateAwareTable { (stateCatalog, tableName) => + stateCatalog.resetLoadTableCalls() + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + RelationCache.empty) + val options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "5")) + val read = UnresolvedRelation(tableName.split("\\.").toSeq, options) + val write = read.requireWritePrivileges(Set(TableWritePrivilege.INSERT)) + + def resolve(relation: UnresolvedRelation): DataSourceV2Relation = { + resolver.resolveRelation(relation).flatMap(_.collectFirst { + case r: DataSourceV2Relation => r + }).getOrElse(fail(s"failed to resolve ${relation.name} as a v2 relation")) + } + + AnalysisContext.withNewAnalysisContext { + val readRelation = resolve(read) + val writeRelation = resolve(write) + + assert(readRelation.table ne writeRelation.table) + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 1) + } + + assert(stateCatalog.loadTableCalls.size == 2) + assert(stateCatalog.loadTableCalls.count(_._1.writePrivileges().isEmpty) == 1) + assert(stateCatalog.loadTableCalls.count( + _._1.writePrivileges().contains(TableWritePrivilege.INSERT)) == 1) + assert(stateCatalog.loadTableCalls.forall(_._2.get("snapshot") == "s1")) + assert(stateCatalog.loadTableCalls.forall(_._2.get("split-size") == "5")) + } + } + + test("transaction V2TableReference skips shared lookup and writes bypass query caches") { + withStateAwareTable { (stateCatalog, tableName) => + val original = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + .queryExecution + .analyzed + .collectFirst { case r: DataSourceV2Relation => r } + .getOrElse(fail("expected a v2 relation")) + val readRef = V2TableReference.createForTransaction(original) + val otherReadRef = V2TableReference.createForTransaction(original.copy( + options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "9")))) + val writeRef = V2TableReference.createForWriteTarget(original) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = (_, _) => { + sharedRelationCacheLookups += 1 + None + } + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + sharedRelationCache) + + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + AnalysisContext.withNewAnalysisContext { + val readRelation = resolver.resolveReference(readRef).asInstanceOf[DataSourceV2Relation] + val cachedReadRelation = + resolver.resolveReference(readRef).asInstanceOf[DataSourceV2Relation] + val otherReadRelation = + resolver.resolveReference(otherReadRef).asInstanceOf[DataSourceV2Relation] + val writeRelation = resolver.resolveReference(writeRef).asInstanceOf[DataSourceV2Relation] + val readAfterWriteRelation = + resolver.resolveReference(readRef).asInstanceOf[DataSourceV2Relation] + + assert(readRelation.table eq cachedReadRelation.table) + assert(readRelation.table eq otherReadRelation.table) + assert(readRelation.table eq readAfterWriteRelation.table) + assert(readRelation.table ne writeRelation.table) + assert(readRelation.options.get("split-size") == "5") + assert(otherReadRelation.options.get("split-size") == "9") + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + } + + assert(sharedRelationCacheLookups == 0) + assert(stateCatalog.singleArgLoads.get() == 2) + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._2.get("snapshot") == "s1") + assert(stateCatalog.loadTableCalls.head._2.get("split-size") == "5") + } + } + + test("nested view resolution shares the query table-state cache") { + withStateAwareTable { (stateCatalog, tableName) => + withView("state_nested_view") { + sql(s"CREATE VIEW state_nested_view AS SELECT * FROM $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 5)") + stateCatalog.resetLoadTableCalls() + + val df = sql(s"SELECT v.id FROM state_nested_view v JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON v.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(relations.head.table eq relations.last.table) + val stateLoads = stateCatalog.loadTableCalls.count(_._2.get("snapshot") == "s1") + assert(stateLoads == 1, s"expected one nested-view table load, got: $stateLoads") + } + } + } + + test("temporary-view V2TableReference consults shared cache only for the initial pin") { + withStateAwareTable { (stateCatalog, tableName) => + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + .queryExecution + .analyzed + .collectFirst { case r: DataSourceV2Relation => r } + .getOrElse(fail("expected a v2 relation")) + val initialRef = V2TableReference.createForTempView(cached, Seq("state_view")) + val otherOptionsRef = V2TableReference.createForTempView( + cached.copy(options = new CaseInsensitiveStringMap( + java.util.Map.of("snapshot", "s1", "split-size", "9"))), + Seq("state_view")) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = (_, _) => { + sharedRelationCacheLookups += 1 + Some(cached) + } + val resolver = new RelationResolution( + spark.sessionState.catalogManager, + sharedRelationCache) + + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + AnalysisContext.withNewAnalysisContext { + val initialRelation = + resolver.resolveReference(initialRef).asInstanceOf[DataSourceV2Relation] + val cachedRelation = + resolver.resolveReference(initialRef).asInstanceOf[DataSourceV2Relation] + val otherOptionsRelation = + resolver.resolveReference(otherOptionsRef).asInstanceOf[DataSourceV2Relation] + + assert(initialRelation.table eq cached.table) + assert(cachedRelation.table eq cached.table) + assert(otherOptionsRelation.table eq cached.table) + assert(initialRelation.options.get("split-size") == "5") + assert(otherOptionsRelation.options.get("split-size") == "9") + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + } + + assert(sharedRelationCacheLookups == 1) + assert(stateCatalog.singleArgLoads.get() == 1) + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._2.get("snapshot") == "s1") + assert(stateCatalog.loadTableCalls.head._2.get("split-size") == "5") + } + } + + test("temporary-view re-resolution preserves the CacheManager table pin") { + withStateAwareTable { (stateCatalog, tableName) => + withTempView("state_view") { + val cached = spark.read + .option("snapshot", "s1") + .option("split-size", "5") + .table(tableName) + cached.cache() + try { + cached.collect() + val cachedTable = cached.queryExecution.analyzed.collectFirst { + case r: DataSourceV2Relation => r.table + }.getOrElse(fail("expected a cached v2 relation")) + cached.createOrReplaceTempView("state_view") + stateCatalog.resetLoadTableCalls() + stateCatalog.singleArgLoads.set(0) + + val df = sql(s"SELECT v.id FROM state_view v JOIN $tableName " + + s"WITH (`snapshot` = 's1', `split-size` = 9) b ON v.id = b.id") + val relations = df.queryExecution.analyzed.collect { + case r: DataSourceV2Relation if r.options.containsKey("snapshot") => r + } + + assert(relations.size == 2) + assert(relations.map(_.options.get("split-size")).sorted == Seq("5", "9")) + assert(relations.forall(_.table eq cachedTable)) + assert(stateCatalog.singleArgLoads.get() == 1, + s"expected one table load while re-resolving the query, got: " + + stateCatalog.singleArgLoads.get()) + assert(stateCatalog.loadTableCalls.size == 1) + assert(stateCatalog.loadTableCalls.head._2.get("snapshot") == "s1") + assert(stateCatalog.loadTableCalls.head._2.get("split-size") == "5") + } finally { + cached.unpersist() + } + } + } + } + test("SPARK-58389: repeated references with the same options load the table once") { val t1 = s"${catalogAndNamespace}table" withTable(t1) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala index ee3eaa789ffc4..23ec3f48bbbed 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/command/PlanResolutionSuite.scala @@ -21,7 +21,7 @@ import java.net.URI import java.util.Collections import org.mockito.ArgumentMatchers.any -import org.mockito.Mockito.{mock, when} +import org.mockito.Mockito.{mock, when, withSettings} import org.mockito.invocation.InvocationOnMock import org.apache.spark.SparkUnsupportedOperationException @@ -38,6 +38,7 @@ import org.apache.spark.sql.catalyst.util.TypeUtils.toSQLId import org.apache.spark.sql.connector.FakeV2Provider import org.apache.spark.sql.connector.catalog.{CatalogManager, Column, ColumnDefaultValue, Identifier, SupportsDelete, Table, TableCapability, TableCatalog, TableChange, TableContext, TableWritePrivilege, V1Table} import org.apache.spark.sql.connector.catalog.CatalogManager.SESSION_CATALOG_NAME +import org.apache.spark.sql.connector.catalog.SupportsTableStateOptions import org.apache.spark.sql.connector.expressions.{LiteralValue, Transform} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.sql.execution.datasources.{CreateTable => CreateTableV1} @@ -184,7 +185,7 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { val ident = invocation.getArguments()(0).asInstanceOf[Identifier] val version = invocation.getArguments()(1).asInstanceOf[String] (ident.name, version) match { - case ("tab", "v1") => table + case ("tab", "v1" | "v2") => table case ("tab", _) => throw new RuntimeException("Unknown version: " + version) case _ => throw new NoSuchTableException(Seq(ident.name)) } @@ -3392,6 +3393,7 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { AnalysisContext.withNewAnalysisContext { val ctx = AnalysisContext.get assert(ctx.relationCache.isEmpty) + assert(ctx.tableCache.isEmpty) // create two unresolved relations without time travel val unresolved1 = UnresolvedRelation(Seq("testcat", "tab")) @@ -3408,6 +3410,7 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { // after first resolution, cache should have 1 entry (without time travel) assert(ctx.relationCache.size == 1) assert(ctx.relationCache.keys.head.timeTravelSpec.isEmpty) + assert(ctx.tableCache.size == 1) // create unresolved relation with time travel spec val timeTravelSpec = AsOfVersion("v1") @@ -3422,6 +3425,17 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { // after time travel resolution, cache should have 2 entries (with and without time travel) assert(ctx.relationCache.size == 2) + assert(ctx.tableCache.size == 2) + + val otherTimeTravelSpec = AsOfVersion("v2") + val resolved4 = resolve( + UnresolvedRelation(Seq("testcat", "tab")), + Some(otherTimeTravelSpec)) + assert(resolved4.timeTravelSpec.contains(otherTimeTravelSpec)) + + // Distinct parsed time-travel specs are distinct state pins even with identical raw options. + assert(ctx.relationCache.size == 3) + assert(ctx.tableCache.size == 3) } } @@ -3483,8 +3497,8 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { cacheOpts: java.util.Map[String, String], readOpts: java.util.Map[String, String]): DataSourceV2Relation = { AnalysisContext.withNewAnalysisContext { - val sharedCache: RelationCache = (_, _) => Some(cachedRelationWith(cacheOpts)) - val rule = new RelationResolution(catalogManagerWithDefault, sharedCache) + val sharedRelationCache: RelationCache = (_, _) => Some(cachedRelationWith(cacheOpts)) + val rule = new RelationResolution(catalogManagerWithDefault, sharedRelationCache) val unresolved = UnresolvedRelation(Seq("testcat", "tab"), new CaseInsensitiveStringMap(readOpts)) rule.resolveRelation(unresolved) match { @@ -3513,6 +3527,124 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { "differing options should freshly load, not reuse the cached relation") } + test("table-state cache consults shared relation cache only while establishing the pin") { + def newTable(id: String): Table = { + val t = mock(classOf[Table]) + when(t.id()).thenReturn(id) + when(t.name()).thenReturn("tab") + when(t.columns()).thenReturn(Array(Column.create("i", IntegerType))) + when(t.capabilities()).thenReturn(java.util.Set.of()) + t + } + + def options(splitSize: String): CaseInsensitiveStringMap = { + new CaseInsensitiveStringMap( + java.util.Map.of("state", "s1", "split-size", splitSize)) + } + + def run( + firstSplitSize: String, + secondSplitSize: String, + sharedRelationCacheEntry: ( + TableCatalog, + Identifier, + Table, + Table) => Option[LogicalPlan]): ( + DataSourceV2Relation, + DataSourceV2Relation, + Table, + Table, + Int, + Int) = { + val currentTable = newTable("table-id") + val cachedTable = newTable("table-id") + val catalogBase = mock( + classOf[TableCatalog], + withSettings().extraInterfaces(classOf[SupportsTableStateOptions])) + val catalog = catalogBase + .asInstanceOf[TableCatalog with SupportsTableStateOptions] + when(catalog.name()).thenReturn("statecat") + when(catalog.tableStateOptionKeys()).thenReturn(java.util.Set.of("state")) + var loads = 0 + when(catalog.loadTable( + any[Identifier], + any[TableContext], + any[CaseInsensitiveStringMap])).thenAnswer((_: InvocationOnMock) => { + loads += 1 + currentTable + }) + + val manager = mock(classOf[CatalogManager]) + when(manager.catalog(any())).thenReturn(catalog) + when(manager.v1SessionCatalog).thenReturn(v1SessionCatalog) + val ident = Identifier.of(Array.empty[String], "tab") + val sharedRelationCacheCandidate = + sharedRelationCacheEntry(catalog, ident, currentTable, cachedTable) + var sharedRelationCacheLookups = 0 + val sharedRelationCache: RelationCache = + (_, _) => { + sharedRelationCacheLookups += 1 + sharedRelationCacheCandidate + } + val resolver = new RelationResolution(manager, sharedRelationCache) + + def resolveWith(splitSize: String): DataSourceV2Relation = { + val unresolved = UnresolvedRelation(Seq("statecat", "tab"), options(splitSize)) + resolver.resolveRelation(unresolved) match { + case Some(AsDataSourceV2Relation(relation)) => relation + case other => fail(s"failed to resolve as v2 relation: $other") + } + } + + AnalysisContext.withNewAnalysisContext { + val first = resolveWith(firstSplitSize) + val second = resolveWith(secondSplitSize) + assert(AnalysisContext.get.tableCache.size == 1) + assert(AnalysisContext.get.relationCache.size == 2) + (first, second, currentTable, cachedTable, loads, sharedRelationCacheLookups) + } + } + + def cachedRelation( + catalog: TableCatalog, + ident: Identifier, + table: Table, + splitSize: String, + tag: Long): DataSourceV2Relation = { + val relation = DataSourceV2Relation.create( + table, + Some(catalog), + Some(ident), + options(splitSize)) + relation.setTagValue(LogicalPlan.PLAN_ID_TAG, tag) + relation + } + + // Situation A: the full-option cached match establishes the initial pin. The later same-state + // lookup reuses that pin without consulting the shared relation cache. + val situationA = run("5", "9", { (catalog, ident, _, cachedTable) => + Some(cachedRelation(catalog, ident, cachedTable, "5", 2L)) + }) + assert(situationA._1.table eq situationA._4) + assert(situationA._2.table eq situationA._4) + assert(situationA._1.getTagValue(LogicalPlan.PLAN_ID_TAG).contains(2L)) + assert(situationA._2.getTagValue(LogicalPlan.PLAN_ID_TAG).isEmpty) + assert(situationA._5 == 1) + assert(situationA._6 == 1) + + // Situation B: a nonmatching option bag resolves first and pins the current Table. The later + // table-cache hit does not consult the shared relation cache, so its full-option entry cannot + // affect the concrete pin. + val situationB = run("9", "5", { (catalog, ident, _, cachedTable) => + Some(cachedRelation(catalog, ident, cachedTable, "5", 6L)) + }) + assert(situationB._1.table eq situationB._3) + assert(situationB._2.table eq situationB._3) + assert(situationB._2.getTagValue(LogicalPlan.PLAN_ID_TAG).isEmpty) + assert(situationB._5 == 1) + assert(situationB._6 == 1) + } + private def resolve( unresolvedRelation: UnresolvedRelation, timeTravelSpec: Option[TimeTravelSpec] = None,