diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/util/CaseInsensitiveStringMap.java b/sql/catalyst/src/main/java/org/apache/spark/sql/util/CaseInsensitiveStringMap.java index 0c85f9b41cbaf..3a0de51c96b00 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/util/CaseInsensitiveStringMap.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/util/CaseInsensitiveStringMap.java @@ -39,7 +39,7 @@ * This is used to pass options to v2 implementations to ensure consistent case insensitivity. *

* Methods that return keys in this map, like {@link #entrySet()} and {@link #keySet()}, return - * keys converted to lower case. This map doesn't allow null key. + * keys converted to lower case. This map is read-only and does not allow null keys. * * @since 3.0.0 */ @@ -58,15 +58,16 @@ public static CaseInsensitiveStringMap empty() { public CaseInsensitiveStringMap(Map originalMap) { original = new HashMap<>(originalMap); - delegate = new HashMap<>(originalMap.size()); + Map normalizedMap = new HashMap<>(originalMap.size()); for (Map.Entry entry : originalMap.entrySet()) { String key = toLowerCase(entry.getKey()); - if (delegate.containsKey(key)) { + if (normalizedMap.containsKey(key)) { logger.warn("Converting duplicated key {} into CaseInsensitiveStringMap.", MDC.of(LogKeys.KEY, entry.getKey())); } - delegate.put(key, entry.getValue()); + normalizedMap.put(key, entry.getValue()); } + delegate = Collections.unmodifiableMap(normalizedMap); } @Override 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 fcbec9db78c9e..89e5db467c9ee 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 @@ -286,8 +286,8 @@ class RelationResolution( // time-travel and write-privilege loads cannot. // // 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. + // mixin): `CatalogV2Util.loadTableWithStateOptions` would call `asTableCatalog` and + // throw MISSING_CATALOG_ABILITY.TABLES, masking the legitimate view-resolution path. val relation: Option[Relation] = pinnedTable.orElse { catalog match { case mc: RelationCatalog @@ -302,12 +302,12 @@ class RelationResolution( val tableSide: Option[Table] = if ( CatalogV2Util.isSessionCatalog(catalog) || catalog.isInstanceOf[TableCatalog] ) { - CatalogV2Util.loadTable( + CatalogV2Util.loadTableWithStateOptions( catalog, ident, + tableKey.stateOptions, finalTimeTravelSpec, - Option(writePrivileges), - finalOptions) + Option(writePrivileges)) } else { None } @@ -406,8 +406,9 @@ class RelationResolution( catalog: CatalogPlugin, ident: Identifier, table: Table, - options: CaseInsensitiveStringMap): Option[DataSourceV2Relation] = { - CatalogV2Util.lookupCachedRelation(sharedRelationCache, catalog, ident, table, options, conf) + stateOptions: CaseInsensitiveStringMap): Option[DataSourceV2Relation] = { + CatalogV2Util.lookupCachedRelationWithStateOptions( + sharedRelationCache, catalog, ident, table, stateOptions, conf) } private def adaptCachedRelation(cached: LogicalPlan, planId: Option[Long]): LogicalPlan = { @@ -534,7 +535,8 @@ class RelationResolution( case Some(pinnedTable) => createRelation(ref, catalog, pinnedTable) case None => - val table = CatalogV2Util.getTable(catalog, ref.identifier, options = ref.options) + val table = CatalogV2Util.getTableWithStateOptions( + catalog, ref.identifier, tableKey.stateOptions) val sharedCacheMatch = if (ref.context.sharedCacheable) { lookupSharedRelationCache(catalog, ref.identifier, table, tableKey.stateOptions) } else { 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 9660ede921214..39dc1b853e05b 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 @@ -475,9 +475,24 @@ private[sql] object CatalogV2Util { ident: Identifier, timeTravelSpec: Option[TimeTravelSpec] = None, writePrivilegesString: Option[String] = None, - options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()): Option[Table] = + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()): Option[Table] = { + val stateOptions = extractTableStateOptions(catalog, options) + loadTableWithStateOptions( + catalog, ident, stateOptions, timeTravelSpec, writePrivilegesString) + } + + /** + * Loads a table using table-state options already projected by the caller. + */ + def loadTableWithStateOptions( + catalog: CatalogPlugin, + ident: Identifier, + stateOptions: CaseInsensitiveStringMap, + timeTravelSpec: Option[TimeTravelSpec] = None, + writePrivilegesString: Option[String] = None): Option[Table] = try { - Option(getTable(catalog, ident, timeTravelSpec, writePrivilegesString, options)) + Option(getTableWithStateOptions( + catalog, ident, stateOptions, timeTravelSpec, writePrivilegesString)) } catch { case _: NoSuchTableException => None case _: NoSuchDatabaseException => None @@ -510,13 +525,26 @@ private[sql] object CatalogV2Util { timeTravelSpec: Option[TimeTravelSpec] = None, writePrivilegesString: Option[String] = None, options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()): Table = { + val stateOptions = extractTableStateOptions(catalog, options) + getTableWithStateOptions( + catalog, ident, stateOptions, timeTravelSpec, writePrivilegesString) + } + + /** + * Loads a table using table-state options already projected by the caller. + */ + def getTableWithStateOptions( + catalog: CatalogPlugin, + ident: Identifier, + stateOptions: CaseInsensitiveStringMap, + timeTravelSpec: Option[TimeTravelSpec] = None, + writePrivilegesString: Option[String] = None): Table = { val timeTravel: TimeTravel = timeTravelSpec match { case Some(v: AsOfVersion) => new TimeTravel.AsOfVersion(v.version) case Some(ts: AsOfTimestamp) => new TimeTravel.AsOfTimestamp(ts.timestamp) case None => null } val context = new TableContext(timeTravel, parseWritePrivileges(writePrivilegesString)) - val stateOptions = extractTableStateOptions(catalog, options) catalog.asTableCatalog.loadTable(ident, context, stateOptions) } @@ -604,18 +632,22 @@ private[sql] object CatalogV2Util { loadTable(catalog, ident).map(DataSourceV2Relation.create(_, Some(catalog), Some(ident))) } - def lookupCachedRelation( + /** + * Looks up a cached relation using table-state options already projected by the caller. + * Reusing the projection keeps table pinning and shared-cache lookup on the same state key. + */ + def lookupCachedRelationWithStateOptions( cache: RelationCache, catalog: CatalogPlugin, ident: Identifier, table: Table, - options: CaseInsensitiveStringMap, + stateOptions: CaseInsensitiveStringMap, conf: SQLConf): Option[DataSourceV2Relation] = { cache.lookup( catalog, ident, Some(table.id), - extractTableStateOptions(catalog, options), + stateOptions, conf.resolver).collect { case r: DataSourceV2Relation => r } } 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 e128adebdff81..7c2de5e698906 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 @@ -17,14 +17,17 @@ package org.apache.spark.sql.connector.catalog +import scala.collection.mutable + import org.mockito.ArgumentCaptor import org.mockito.ArgumentMatchers.{any, eq => mockEq} import org.mockito.Mockito.{mock, verify, when} +import org.mockito.invocation.InvocationOnMock import org.apache.spark.{SparkFunSuite, SparkIllegalArgumentException} import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.analysis.{ - AsOfTimestamp, AsOfVersion, TimeTravelSpec, UnresolvedRelation} + AsOfTimestamp, AsOfVersion, TableCacheKey, TimeTravelSpec, UnresolvedRelation} import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{IntegerType, StructType} @@ -197,6 +200,68 @@ class CatalogV2UtilSuite extends SparkFunSuite { mockEq(expected)) } + test("getTableWithStateOptions rejects catalog mutation of a retained cache key") { + val catalog = mock(classOf[TableCatalog]) + val ident = Identifier.of(Array("ns"), "table") + val table = mock(classOf[Table]) + var catalogOptions: CaseInsensitiveStringMap = null + var mutationBlocked = false + when(catalog.loadTable(any[Identifier], any[TableContext], any[CaseInsensitiveStringMap])) + .thenAnswer((invocation: InvocationOnMock) => { + catalogOptions = invocation.getArgument[CaseInsensitiveStringMap](2) + intercept[UnsupportedOperationException] { + catalogOptions.entrySet().iterator().next().setValue("s2") + } + mutationBlocked = true + table + }) + val stateOptions = + new CaseInsensitiveStringMap(java.util.Map.of("SnApShOt", "s1")) + val retainedKey = TableCacheKey(catalog, ident, None, stateOptions) + + val loaded = CatalogV2Util.getTableWithStateOptions(catalog, ident, stateOptions) + val tableCache = mutable.Map(retainedKey -> loaded) + val sameStateKey = TableCacheKey( + catalog, + ident, + None, + new CaseInsensitiveStringMap(java.util.Map.of("snapshot", "s1"))) + val differentStateKey = TableCacheKey( + catalog, + ident, + None, + new CaseInsensitiveStringMap(java.util.Map.of("snapshot", "s2"))) + + assert(mutationBlocked, "the catalog fixture did not attempt to mutate its map") + assert(catalogOptions.get("snapshot") == "s1") + assert(catalogOptions.asCaseSensitiveMap().containsKey("SnApShOt")) + assert(tableCache.get(sameStateKey).contains(table)) + assert(!tableCache.contains(differentStateKey)) + } + + test("getTableWithStateOptions preserves effective values for duplicate-case keys") { + val catalog = mock(classOf[TableCatalog]) + val ident = Identifier.of(Array("ns"), "table") + val table = mock(classOf[Table]) + var catalogOptions: CaseInsensitiveStringMap = null + when(catalog.loadTable(any[Identifier], any[TableContext], any[CaseInsensitiveStringMap])) + .thenAnswer((invocation: InvocationOnMock) => { + catalogOptions = invocation.getArgument[CaseInsensitiveStringMap](2) + table + }) + val options = new java.util.LinkedHashMap[String, String]() + options.put("snapshot", "lower-value") + options.put("SNAPSHOT", "upper-value") + val stateOptions = new CaseInsensitiveStringMap(options) + + CatalogV2Util.getTableWithStateOptions(catalog, ident, stateOptions) + + assert(stateOptions.asCaseSensitiveMap().size() == 2) + assert(stateOptions.get("snapshot") == "upper-value") + assert(catalogOptions.get("snapshot") == stateOptions.get("snapshot")) + assert(catalogOptions.asCaseSensitiveMap() == stateOptions.asCaseSensitiveMap()) + } + test("getTable forwards no options when a catalog declares no table-state options") { val catalog = mock(classOf[TableCatalog]) when(catalog.tableStateOptionKeys()).thenCallRealMethod() diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/util/CaseInsensitiveStringMapSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/util/CaseInsensitiveStringMapSuite.scala index 932fb0a733371..fb5e39861fb52 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/util/CaseInsensitiveStringMapSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/util/CaseInsensitiveStringMapSuite.scala @@ -45,6 +45,31 @@ class CaseInsensitiveStringMapSuite extends SparkFunSuite { assert(options.values().asScala.toSeq == Seq("valUE")) } + test("collection views are read-only") { + def checkMutationBlocked(mutate: CaseInsensitiveStringMap => Any): Unit = { + val options = new CaseInsensitiveStringMap(Map("kEy" -> "valUE").asJava) + + intercept[UnsupportedOperationException] { + mutate(options) + } + + assert(options.get("key") == "valUE") + assert(options.asCaseSensitiveMap().asScala == Map("kEy" -> "valUE")) + } + + checkMutationBlocked(_.keySet().remove("key")) + checkMutationBlocked(_.values().remove("valUE")) + checkMutationBlocked(_.entrySet().iterator().next().setValue("newValue")) + checkMutationBlocked { options => + val iterator = options.entrySet().iterator() + iterator.next() + iterator.remove() + } + checkMutationBlocked { options => + options.replaceAll((_, _) => "newValue") + } + } + test("getInt") { val options = new CaseInsensitiveStringMap(Map("numFOo" -> "1", "foo" -> "bar").asJava) assert(options.getInt("numFOO", 10) == 1) @@ -107,5 +132,9 @@ class CaseInsensitiveStringMapSuite extends SparkFunSuite { intercept[UnsupportedOperationException] { caseSensitiveMap.put("kEy", "valUE") } + intercept[UnsupportedOperationException] { + caseSensitiveMap.entrySet().iterator().next().setValue("newValue") + } + assert(caseSensitiveMap.equals(originalMap)) } } 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 6653a10c87c1f..da5f930795e70 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 @@ -91,13 +91,13 @@ private[sql] object V2TableRefreshUtil extends SQLConfHelper with Logging { val stateOptions = CatalogV2Util.extractTableStateOptions(catalog, r.options) val currentTable = currentTables.getOrElseUpdate((catalog, ident, stateOptions), { val tableName = V2TableUtil.toQualifiedName(catalog, ident) - lookupCachedRelation(spark, catalog, ident, r.table, r.options) match { + lookupCachedRelation(spark, catalog, ident, r.table, stateOptions) match { case Some(cached) => logDebug(s"Refreshing table metadata for $tableName using shared relation cache") cached.table case _ => logDebug(s"Refreshing table metadata for $tableName using catalog") - CatalogV2Util.getTable(catalog, ident, options = r.options) + CatalogV2Util.getTableWithStateOptions(catalog, ident, stateOptions) } }) validateTableIdentity(currentTable, r) @@ -112,9 +112,9 @@ private[sql] object V2TableRefreshUtil extends SQLConfHelper with Logging { catalog: TableCatalog, ident: Identifier, table: Table, - options: CaseInsensitiveStringMap): Option[DataSourceV2Relation] = { - CatalogV2Util.lookupCachedRelation( - spark.sharedState.relationCache, catalog, ident, table, options, conf) + stateOptions: CaseInsensitiveStringMap): Option[DataSourceV2Relation] = { + CatalogV2Util.lookupCachedRelationWithStateOptions( + spark.sharedState.relationCache, catalog, ident, table, stateOptions, conf) } // it is not safe to allow any schema changes in commands (e.g. CTAS, RTAS, MERGE) 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 0c9032fbeb001..8b7fd2bb80e3a 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 @@ -71,10 +71,17 @@ class LoadCountingInMemoryCatalog extends InMemoryCatalog { } class StateAwareInMemoryCatalog extends LoadCountingInMemoryCatalog { + private val _tableStateOptionKeyCalls = new AtomicInteger(0) + + def tableStateOptionKeyCalls: Int = _tableStateOptionKeyCalls.get() + def resetTableStateOptionKeyCalls(): Unit = _tableStateOptionKeyCalls.set(0) + // Include Spark's internal marker so the write-context test detects if it leaks before state // option projection. Production catalogs should declare only raw user option keys. - override def tableStateOptionKeys(): util.Set[String] = + override def tableStateOptionKeys(): util.Set[String] = { + _tableStateOptionKeyCalls.incrementAndGet() util.Set.of("snapshot", UnresolvedRelation.REQUIRED_WRITE_PRIVILEGES) + } } class NullReturningInMemoryCatalog extends InMemoryCatalog { @@ -142,8 +149,12 @@ class StateAwareV2InMemoryRelationCatalog extends V2InMemoryRelationCatalog { private val _loadTableCalls = mutable.ArrayBuffer.empty[(Identifier, TableContext, CaseInsensitiveStringMap)] private val _loadViewCalls = mutable.ArrayBuffer.empty[Identifier] + private val _tableStateOptionKeyCalls = new AtomicInteger(0) - override def tableStateOptionKeys(): util.Set[String] = util.Set.of("snapshot") + override def tableStateOptionKeys(): util.Set[String] = { + _tableStateOptionKeyCalls.incrementAndGet() + util.Set.of("snapshot") + } override def loadTable( ident: Identifier, @@ -168,11 +179,13 @@ class StateAwareV2InMemoryRelationCatalog extends V2InMemoryRelationCatalog { def loadTableCalls: Seq[(Identifier, TableContext, CaseInsensitiveStringMap)] = _loadTableCalls.toSeq def loadViewCalls: Seq[Identifier] = _loadViewCalls.toSeq + def tableStateOptionKeyCalls: Int = _tableStateOptionKeyCalls.get() def resetDispatchCalls(): Unit = { resetLoadRelationCalls() _loadTableCalls.clear() _loadViewCalls.clear() + _tableStateOptionKeyCalls.set(0) } } @@ -1032,13 +1045,24 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { sql(s"CREATE TABLE $t1 (id bigint, data string) USING parquet") relCatalog.resetDispatchCalls() - spark.read + val df = spark.read .option("snapshot", "s1") .option("split-size", "5") .table(t1) - .collect() + df.queryExecution.analyzed + + assertOnlySnapshotTableOptions(relCatalog, "s1", expectedCalls = 1) + assert( + relCatalog.tableStateOptionKeyCalls == 1, + s"expected one analysis projection, got: ${relCatalog.tableStateOptionKeyCalls}") - assertOnlySnapshotTableOptions(relCatalog, "s1", expectedCalls = 2) + relCatalog.resetDispatchCalls() + df.collect() + + assertOnlySnapshotTableOptions(relCatalog, "s1", expectedCalls = 1) + assert( + relCatalog.tableStateOptionKeyCalls == 1, + s"expected one refresh projection, got: ${relCatalog.tableStateOptionKeyCalls}") assert(relCatalog.loadRelationCalls.isEmpty) } } @@ -1703,12 +1727,28 @@ class DataSourceV2OptionSuite extends DatasourceV2SQLBase { stateCatalog.resetLoadTableCalls() stateCatalog.singleArgLoads.set(0) AnalysisContext.withNewAnalysisContext { + stateCatalog.resetTableStateOptionKeyCalls() val initialRelation = resolver.resolveReference(initialRef).asInstanceOf[DataSourceV2Relation] + assert( + stateCatalog.tableStateOptionKeyCalls == 1, + s"expected one initial projection, got: ${stateCatalog.tableStateOptionKeyCalls}") + + stateCatalog.resetTableStateOptionKeyCalls() val cachedRelation = resolver.resolveReference(initialRef).asInstanceOf[DataSourceV2Relation] + assert( + stateCatalog.tableStateOptionKeyCalls == 0, + s"expected no projection on a relation-cache hit, got: " + + stateCatalog.tableStateOptionKeyCalls) + + stateCatalog.resetTableStateOptionKeyCalls() val otherOptionsRelation = resolver.resolveReference(otherOptionsRef).asInstanceOf[DataSourceV2Relation] + assert( + stateCatalog.tableStateOptionKeyCalls == 1, + s"expected one table-cache lookup projection, got: " + + stateCatalog.tableStateOptionKeyCalls) assert(initialRelation.table eq cached.table) assert(cachedRelation.table eq cached.table) 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 1065c5baa5858..f960507b300e5 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 @@ -3583,7 +3583,11 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { val cachedTable = newTable("table-id") val catalog = mock(classOf[TableCatalog]) when(catalog.name()).thenReturn("statecat") - when(catalog.tableStateOptionKeys()).thenReturn(java.util.Set.of("state")) + var tableStateOptionKeyCalls = 0 + when(catalog.tableStateOptionKeys()).thenAnswer((_: InvocationOnMock) => { + tableStateOptionKeyCalls += 1 + java.util.Set.of("state") + }) var loads = 0 when(catalog.loadTable( any[Identifier], @@ -3608,11 +3612,16 @@ class PlanResolutionSuite extends SharedSparkSession with AnalysisTest { val resolver = new RelationResolution(manager, sharedRelationCache) def resolveWith(splitSize: String): DataSourceV2Relation = { + tableStateOptionKeyCalls = 0 val unresolved = UnresolvedRelation(Seq("statecat", "tab"), options(splitSize)) - resolver.resolveRelation(unresolved) match { + val resolved = resolver.resolveRelation(unresolved) match { case Some(AsDataSourceV2Relation(relation)) => relation case other => fail(s"failed to resolve as v2 relation: $other") } + assert( + tableStateOptionKeyCalls == 1, + s"expected one table-state option projection, got: $tableStateOptionKeyCalls") + resolved } AnalysisContext.withNewAnalysisContext {