Introduce PlanCacheSchemaKey and multi-schema QueryCacheKey (cross-schema 1/6)

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizer.java

@@ -36,8 +36,10 @@
 import com.apple.foundationdb.relational.generated.RelationalParserBaseVisitor;
 import com.apple.foundationdb.relational.recordlayer.metadata.DataTypeUtils;
 import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerInvokedRoutine;
 import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerSchemaTemplate;
+import com.apple.foundationdb.relational.recordlayer.query.cache.PlanCacheSchemaKey;
 import com.apple.foundationdb.relational.recordlayer.query.cache.QueryCacheKey;
+import com.google.common.collect.ImmutableSortedMap;
 import com.apple.foundationdb.relational.recordlayer.util.ExceptionUtil;
 import com.apple.foundationdb.relational.util.Assert;
 import com.google.common.annotations.VisibleForTesting;
@@ -630,10 +632,10 @@
             }
         }
         return new NormalizationResult(
-                recordLayerSchemaTemplate.getName(),
+                PlanCacheSchemaKey.of(recordLayerSchemaTemplate.getName()),
                 QueryCacheKey.of(astNormalizer.getCanonicalSqlString(), getQuerySpecificPlannerConfig(plannerConfiguration, astNormalizer.getQueryOptions()),
                         recordLayerSchemaTemplate.getTransactionBoundMetadataAsString(),
-                        recordLayerSchemaTemplate.getVersion(), userVersion),
+                        ImmutableSortedMap.of(recordLayerSchemaTemplate.getName(), recordLayerSchemaTemplate.getVersion()), userVersion),
                 astNormalizer.getQueryExecutionParameters(),
                 parseTreeInfo.getRootContext(),
                 astNormalizer.getQueryCachingFlags(),
@@ -677,7 +679,7 @@
         }
 
         @Nonnull
-        private final String schemaTemplateName;
+        private final PlanCacheSchemaKey schemaKey;
 
         @Nonnull
         private final QueryCacheKey queryCacheKey;
@@ -697,14 +699,14 @@
         @Nonnull
         private final String query;
 
-        public NormalizationResult(@Nonnull final String schemaTemplateName,
+        public NormalizationResult(@Nonnull final PlanCacheSchemaKey schemaKey,
                                    @Nonnull final QueryCacheKey queryCacheKey,
                                    @Nonnull final QueryExecutionContext queryExecutionContext,
                                    @Nonnull final ParseTree parseTree,
                                    @Nonnull final Set<QueryCachingFlags> queryCachingFlags,
                                    @Nonnull final Options queryOptions,
                                    @Nonnull final String query) {
-            this.schemaTemplateName = schemaTemplateName;
+            this.schemaKey = schemaKey;
             this.queryCacheKey = queryCacheKey;
             this.queryExecutionContext = queryExecutionContext;
             this.parseTree = parseTree;
@@ -714,8 +716,8 @@
         }
 
         @Nonnull
-        public String getSchemaTemplateName() {
-            return schemaTemplateName;
+        public PlanCacheSchemaKey getSchemaKey() {
+            return schemaKey;
         }
 
         @Nonnull

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/PlanGenerator.java

@@ -172,8 +172,8 @@
             // otherwise, lookup the query in the cache
             final var planEquivalence = PhysicalPlanEquivalence.of(astHashResult.getQueryExecutionContext().getEvaluationContext());
             return planContext.getMetricsCollector().clock(RelationalMetric.RelationalEvent.CACHE_LOOKUP, () ->
                     cache.get().reduce(
-                            astHashResult.getSchemaTemplateName(),
+                            astHashResult.getSchemaKey(),
                             astHashResult.getQueryCacheKey(),
                             planEquivalence,
                             () -> {

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/PlanCacheSchemaKey.java

@@ -0,0 +1,84 @@
+/*
+ * PlanCacheSchemaKey.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2021-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed 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 com.apple.foundationdb.relational.recordlayer.query.cache;
+
+import com.apple.foundationdb.annotation.API;
+import com.google.common.collect.ImmutableSortedSet;
+
+import javax.annotation.Nonnull;
+import java.util.Collection;
+
+/**
+ * Primary key for {@link RelationalPlanCache}. Represents the set of schema template names that a
+ * cached plan was compiled against. Single-schema queries produce a one-element set; cross-schema
+ * queries produce a multi-element set.
+ * <p>
+ * The set is always sorted so that {@code {"a","b"}} and {@code {"b","a"}} map to the same bucket.
+ */
+@API(API.Status.EXPERIMENTAL)
+public final class PlanCacheSchemaKey {
+
+    @Nonnull
+    private final ImmutableSortedSet<String> schemaNames;
+
+    private final int memoizedHashCode;
+
+    private PlanCacheSchemaKey(@Nonnull final ImmutableSortedSet<String> schemaNames) {
+        this.schemaNames = schemaNames;
+        this.memoizedHashCode = schemaNames.hashCode();
+    }
+
+    @Nonnull
+    public ImmutableSortedSet<String> getSchemaNames() {
+        return schemaNames;
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (this == other) {
+            return true;
+        }
+        if (other == null || getClass() != other.getClass()) {
+            return false;
+        }
+        return schemaNames.equals(((PlanCacheSchemaKey) other).schemaNames);
+    }
+
+    @Override
+    public int hashCode() {
+        return memoizedHashCode;
+    }
+
+    @Override
+    public String toString() {
+        return String.join("|", schemaNames);
+    }
+
+    @Nonnull
+    public static PlanCacheSchemaKey of(@Nonnull final String schemaName) {
+        return new PlanCacheSchemaKey(ImmutableSortedSet.of(schemaName));
+    }
+
+    @Nonnull
+    public static PlanCacheSchemaKey of(@Nonnull final Collection<String> schemaNames) {
+        return new PlanCacheSchemaKey(ImmutableSortedSet.copyOf(schemaNames));
+    }
+}

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/QueryCacheKey.java

@@ -23,7 +23,8 @@
 import com.apple.foundationdb.annotation.API;
 
 import com.apple.foundationdb.relational.recordlayer.query.AstNormalizer;
 import com.apple.foundationdb.relational.recordlayer.query.PlannerConfiguration;
+import com.google.common.collect.ImmutableSortedMap;
 
 import javax.annotation.Nonnull;
 import java.util.Objects;
@@ -32,53 +33,13 @@
  * This is used to look up a plan in the primary cache (see {@link MultiStageCache} for more information).
  * It comprises the following fields:
  * <ul>
- *     <li>The schema template name to which the query is bound. It is necessary to use it so we can segregate
- *     the otherwise identical plans but coming from different schemas (see Example 1 below)</li>
- *     <li>The schema template version, user version, and a bit-set of all readable indexes</li>
- *     <li>The canonical query string where all literals are removed and white spaces are normalised (see example 2)</li>
+ *     <li>The schema template versions to which the query is bound, keyed by schema template name. A plan is
+ *     invalidated when <em>any</em> participating schema bumps its version.</li>
+ *     <li>The user version and a bit-set of all readable indexes.</li>
+ *     <li>The canonical query string where all literals are removed and white spaces are normalised (see example below).</li>
  *     <li>The hash of the query, see {@link AstNormalizer} for more information on how is this generated.</li>
  * </ul>
- * <b>Example1</b>
- * <br>
- * Let us assume we have two schema templates ({@code s1} and {@code s2}) defined as the following:
- * <pre>
- * {@code
- * create schema template s1
- * create table t1(id bigint, col1 bigint, primary key(id))
- *
- * create schema template s2
- * create table t1(id bigint, col1 bigint, col3 bigint, primary key(id))
- * }
- * </pre>
- * If we run a query like this:
- * <pre>
- * {@code
- * create schema /FRL/YOUSSEF/s1s with s1
- * connect: "jdbc:embed:/FRL/YOUSSEF?schema=S1S"
- * select * from t1 where col1 > 42;
- * }
- * </pre>
- * we would compile the query and return a result set comprising two columns {@code id, col1}.
- * we will also cache this query using a {@link QueryCacheKey} key of something like {@code "s1", "select * from t1 where col1 > ? ", 123456789}
- * <br>
- * if we run the <i>same</i> query, however after connecting to a schema that uses a {@code s2} instead:
- * <pre>
- * {@code
- * create schema /FRL/YOUSSEF/s2s with s2
- * connect: "jdbc:embed:/FRL/YOUSSEF?schema=S2S"
- * select * from t1 where col1 > 53;
- * }
- * </pre>
- * we would correctly compile this query and return a result set comprising three columns {@code id, col1, col2}.
- * we will also cache this query using {@link QueryCacheKey} key of something like {@code "s2", "select * from t1 where col1 > ? ", 123456789}
- * <br>
- * Note that without having the schema template name as part of the {@link QueryCacheKey} both keys would be identical.
- * Therefore, we might incorrectly choose the compiled plan of the first query to execute the second query (because we find
- * a match in the cache) causing an error since the result set structure is different because table {@code T1} is defined
- * differently in {@code S1} and {@code S2}.
- * <br>
- * <br>
- * <b>Example 2</b>
+ * <b>Example</b>
  * <br>
  * Although these queries appear different, their canonical representation is the same, and will end up using the same
  * compiled plan.
@@ -108,7 +69,12 @@
     @Nonnull
     private final String auxiliaryMetadata;
 
-    private final int schemaTemplateVersion;
+    /**
+     * Maps each participating schema template name to its version. A cache miss is triggered when
+     * any entry in this map changes, which is the correct behaviour for cross-schema queries.
+     */
+    @Nonnull
+    private final ImmutableSortedMap<String, Integer> schemaVersions;
 
     private final int userVersion;
 
@@ -117,18 +83,15 @@
     private QueryCacheKey(@Nonnull final String canonicalQueryString,
                           @Nonnull final PlannerConfiguration plannerConfiguration,
                           @Nonnull final String auxiliaryMetadata,
-                          int schemaTemplateVersion,
+                          @Nonnull final ImmutableSortedMap<String, Integer> schemaVersions,
                           int userVersion) {
         this.canonicalQueryString = canonicalQueryString;
-        this.schemaTemplateVersion = schemaTemplateVersion;
+        this.schemaVersions = schemaVersions;
         this.userVersion = userVersion;
         this.auxiliaryMetadata = auxiliaryMetadata;
         this.plannerConfiguration = plannerConfiguration;
 
-        // Memoize the hash code. Because this object is used as a key in a hash map, it is important that
-        // hashCode() be quick. Note that this includes information about the query (canonicalQueryString is like the query hash),
-        // the schema template version, and the schema (like the set of readable indexes)
-        this.memoizedHashCode = Objects.hash(canonicalQueryString, schemaTemplateVersion, plannerConfiguration, userVersion, auxiliaryMetadata);
+        this.memoizedHashCode = Objects.hash(canonicalQueryString, schemaVersions, plannerConfiguration, userVersion, auxiliaryMetadata);
     }
 
     @Override
@@ -140,11 +103,11 @@
             return false;
         }
         final var that = (QueryCacheKey) other;
-        return schemaTemplateVersion == that.schemaTemplateVersion &&
-                userVersion == that.userVersion &&
+        return userVersion == that.userVersion &&
                 Objects.equals(canonicalQueryString, that.canonicalQueryString) &&
                 Objects.equals(auxiliaryMetadata, that.auxiliaryMetadata) &&
-                Objects.equals(plannerConfiguration, that.plannerConfiguration);
+                Objects.equals(plannerConfiguration, that.plannerConfiguration) &&
+                Objects.equals(schemaVersions, that.schemaVersions);
     }
 
     @Override
@@ -157,8 +120,9 @@
         return canonicalQueryString;
     }
 
-    public int getSchemaTemplateVersion() {
-        return schemaTemplateVersion;
+    @Nonnull
+    public ImmutableSortedMap<String, Integer> getSchemaVersions() {
+        return schemaVersions;
     }
 
     @Nonnull
@@ -177,16 +141,15 @@
 
     @Override
     public String toString() {
-        return "(" + schemaTemplateVersion + " || " + auxiliaryMetadata + ")" + "||" + canonicalQueryString + "||" + memoizedHashCode;
+        return "(" + schemaVersions + " || " + auxiliaryMetadata + ")" + "||" + canonicalQueryString + "||" + memoizedHashCode;
     }
 
     @Nonnull
     public static QueryCacheKey of(@Nonnull final String query,
                                    @Nonnull final PlannerConfiguration plannerConfiguration,
                                    @Nonnull final String auxiliaryMetadata,
-                                   int schemaTemplateVersion,
+                                   @Nonnull final ImmutableSortedMap<String, Integer> schemaVersions,
                                    int userVersion) {
-        return new QueryCacheKey(query, plannerConfiguration, auxiliaryMetadata, schemaTemplateVersion,
-                userVersion);
+        return new QueryCacheKey(query, plannerConfiguration, auxiliaryMetadata, schemaVersions, userVersion);
     }
 }

fdb-relational-core/src/main/java/com/apple/foundationdb/relational/recordlayer/query/cache/RelationalPlanCache.java

@@ -33,8 +33,8 @@
 /**
  * This is just a specialization of {@link MultiStageCache} with concrete types specific to plan caching.
  */
 @API(API.Status.EXPERIMENTAL)
-public final class RelationalPlanCache extends MultiStageCache<String, QueryCacheKey, PhysicalPlanEquivalence, Plan<?>> {
+public final class RelationalPlanCache extends MultiStageCache<PlanCacheSchemaKey, QueryCacheKey, PhysicalPlanEquivalence, Plan<?>> {
 
     @Nonnull
     private static final TimeUnit DEFAULT_TTL_TIME_UNIT = TimeUnit.MILLISECONDS;
@@ -61,7 +61,7 @@
         super(size, secondarySize, tertiarySize, ttl, ttlTimeUnit, secondaryTtl, secondaryTtlTimeUnit, tertiaryTtl, tertiaryTtlTimeUnit, executor, secondaryExecutor, tertiaryExecutor, ticker);
     }
 
-    public static final class RelationalCacheBuilder extends MultiStageCache.Builder<String, QueryCacheKey, PhysicalPlanEquivalence, Plan<?>, RelationalCacheBuilder> {
+    public static final class RelationalCacheBuilder extends MultiStageCache.Builder<PlanCacheSchemaKey, QueryCacheKey, PhysicalPlanEquivalence, Plan<?>, RelationalCacheBuilder> {
 
         public RelationalCacheBuilder() {
             size = (Integer) (Options.defaultOptions().get(Options.Name.PLAN_CACHE_PRIMARY_MAX_ENTRIES));

fdb-relational-core/src/test/java/com/apple/foundationdb/relational/recordlayer/query/AstNormalizerTests.java

@@ -39,6 +39,7 @@
 import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerSchemaTemplate;
 import com.apple.foundationdb.relational.recordlayer.metadata.RecordLayerTable;
 import com.apple.foundationdb.relational.recordlayer.query.cache.QueryCacheKey;
+import com.google.common.collect.ImmutableSortedMap;
 import com.apple.foundationdb.relational.recordlayer.query.functions.CompiledSqlFunction;
 import com.apple.foundationdb.relational.recordlayer.util.Hex;
 import com.apple.foundationdb.relational.util.Assert;
@@ -1496,8 +1497,9 @@ void visitFullDescribeStatementThrows() throws ReflectiveOperationException {
 
     @Test
     void queryCacheKeyToString() {
-        final var key = QueryCacheKey.of("select ? from testTable", plannerConfiguration, "someAuxiliaryMetadata", 3, 7);
-        final var expected = "(3 || someAuxiliaryMetadata)||select ? from testTable||" + key.hashCode();
+        final var key = QueryCacheKey.of("select ? from testTable", plannerConfiguration, "someAuxiliaryMetadata",
+                ImmutableSortedMap.of("testSchema", 3), 7);
+        final var expected = "({testSchema=3} || someAuxiliaryMetadata)||select ? from testTable||" + key.hashCode();
         Assertions.assertThat(key).hasToString(expected);
     }
 }