Add RecordQueryStoreBindingPlan and cross-schema join support in planner (cross-schema 2/6)

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/EvaluationContext.java

@@ -21,8 +21,10 @@
 package com.apple.foundationdb.record;
 
 import com.apple.foundationdb.annotation.API;
+import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase;
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
 import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
+import com.google.common.collect.ImmutableMap;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -44,7 +46,10 @@ public class EvaluationContext {
     @Nonnull
     private final TypeRepository typeRepository;
 
-    public static final EvaluationContext EMPTY = new EvaluationContext(Bindings.EMPTY_BINDINGS, TypeRepository.EMPTY_SCHEMA);
+    @Nonnull
+    private final ImmutableMap<String, FDBRecordStoreBase<?>> auxiliaryStores;
+
+    public static final EvaluationContext EMPTY = new EvaluationContext(Bindings.EMPTY_BINDINGS, TypeRepository.EMPTY_SCHEMA, ImmutableMap.of());
 
     /**
      * Get an empty evaluation context.
@@ -55,9 +60,11 @@ public static EvaluationContext empty() {
         return EMPTY;
     }
 
-    private EvaluationContext(@Nonnull Bindings bindings, @Nonnull TypeRepository typeRepository) {
+    private EvaluationContext(@Nonnull Bindings bindings, @Nonnull TypeRepository typeRepository,
+                              @Nonnull ImmutableMap<String, FDBRecordStoreBase<?>> auxiliaryStores) {
         this.bindings = bindings;
         this.typeRepository = typeRepository;
+        this.auxiliaryStores = auxiliaryStores;
     }
 
     /**
@@ -68,7 +75,7 @@ private EvaluationContext(@Nonnull Bindings bindings, @Nonnull TypeRepository ty
      */
     @Nonnull
     public static EvaluationContext forBindings(@Nonnull Bindings bindings) {
-        return new EvaluationContext(bindings, TypeRepository.EMPTY_SCHEMA);
+        return new EvaluationContext(bindings, TypeRepository.EMPTY_SCHEMA, ImmutableMap.of());
     }
 
     /**
@@ -80,12 +87,18 @@ public static EvaluationContext forBindings(@Nonnull Bindings bindings) {
      */
     @Nonnull
     public static EvaluationContext forBindingsAndTypeRepository(@Nonnull Bindings bindings, @Nonnull TypeRepository typeRepository) {
-        return new EvaluationContext(bindings, typeRepository);
+        return new EvaluationContext(bindings, typeRepository, ImmutableMap.of());
+    }
+
+    @Nonnull
+    public static EvaluationContext forBindingsAndTypeRepository(@Nonnull Bindings bindings, @Nonnull TypeRepository typeRepository,
+                                                                 @Nonnull ImmutableMap<String, FDBRecordStoreBase<?>> auxiliaryStores) {
+        return new EvaluationContext(bindings, typeRepository, auxiliaryStores);
     }
 
     @Nonnull
     public static EvaluationContext forTypeRepository(@Nonnull TypeRepository typeRepository) {
-        return new EvaluationContext(Bindings.EMPTY_BINDINGS, typeRepository);
+        return new EvaluationContext(Bindings.EMPTY_BINDINGS, typeRepository, ImmutableMap.of());
     }
 
     /**
@@ -97,7 +110,7 @@ public static EvaluationContext forTypeRepository(@Nonnull TypeRepository typeRe
      */
     @Nonnull
     public static EvaluationContext forBinding(@Nonnull String bindingName, @Nullable Object value) {
-        return new EvaluationContext(Bindings.newBuilder().set(bindingName, value).build(), TypeRepository.EMPTY_SCHEMA);
+        return new EvaluationContext(Bindings.newBuilder().set(bindingName, value).build(), TypeRepository.EMPTY_SCHEMA, ImmutableMap.of());
     }
 
     /**
@@ -197,6 +210,31 @@ public TypeRepository getTypeRepository() {
         return typeRepository;
     }
 
+    /**
+     * Returns the auxiliary store bound to the given schema name, or {@code null} if no such
+     * store has been injected. Used by {@code RecordQueryStoreBindingPlan} to redirect execution
+     * to a secondary schema's record store.
+     *
+     * @param schemaName the name of the secondary schema
+     * @return the bound store, or {@code null}
+     */
+    @Nullable
+    public FDBRecordStoreBase<?> getAuxiliaryStore(@Nonnull final String schemaName) {
+        return auxiliaryStores.get(schemaName);
+    }
+
+    /**
+     * Returns a new {@link EvaluationContext} identical to this one except that the given
+     * auxiliary stores are injected. Existing bindings and type repository are preserved.
+     *
+     * @param stores map from schema name to pre-opened record store
+     * @return new context with auxiliary stores
+     */
+    @Nonnull
+    public EvaluationContext withAuxiliaryStores(@Nonnull final ImmutableMap<String, FDBRecordStoreBase<?>> stores) {
+        return new EvaluationContext(bindings, typeRepository, stores);
+    }
+
     /**
      * Construct a builder from this context. This allows the user to create
      * a new <code>EvaluationContext</code> that has all of the same data
@@ -232,7 +270,10 @@ public static EvaluationContextBuilder newBuilder() {
      */
     @Nonnull
     public EvaluationContext withBinding(@Nonnull String bindingName, @Nullable Object value) {
-        return childBuilder().setBinding(bindingName, value).build(typeRepository);
+        return new EvaluationContext(
+                bindings.childBuilder().set(bindingName, value).build(),
+                typeRepository,
+                auxiliaryStores);
     }
 
     /**
@@ -248,6 +289,6 @@ public EvaluationContext withBinding(@Nonnull String bindingName, @Nullable Obje
      * @return a new <code>EvaluationContext</code> with the new binding
      */
     public EvaluationContext withBinding(final Bindings.Internal type, @Nonnull CorrelationIdentifier alias, @Nullable Object value) {
-        return childBuilder().setBinding(type.bindingName(alias.getId()), value).build(typeRepository);
+        return withBinding(type.bindingName(alias.getId()), value);
     }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/EvaluationContextBuilder.java

@@ -21,8 +21,10 @@
 package com.apple.foundationdb.record;
 
 import com.apple.foundationdb.annotation.API;
+import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase;
 import com.apple.foundationdb.record.query.plan.cascades.CorrelationIdentifier;
 import com.apple.foundationdb.record.query.plan.cascades.typing.TypeRepository;
+import com.google.common.collect.ImmutableMap;
 
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -37,12 +39,15 @@
 public class EvaluationContextBuilder {
     @Nonnull
     protected final Bindings.Builder bindings;
+    @Nonnull
+    protected ImmutableMap<String, FDBRecordStoreBase<?>> auxiliaryStores;
 
     /**
      * Create an empty builder.
      */
     protected EvaluationContextBuilder() {
         this.bindings = Bindings.newBuilder();
+        this.auxiliaryStores = ImmutableMap.of();
     }
 
     /**
@@ -54,6 +59,7 @@ protected EvaluationContextBuilder() {
      */
     protected EvaluationContextBuilder(@Nonnull EvaluationContext original) {
         this.bindings = original.getBindings().childBuilder();
+        this.auxiliaryStores = ImmutableMap.of();
     }
 
     /**
@@ -116,6 +122,6 @@ public EvaluationContextBuilder setConstant(@Nonnull CorrelationIdentifier alias
      */
     @Nonnull
     public EvaluationContext build(@Nonnull final TypeRepository typeRepository) {
-        return EvaluationContext.forBindingsAndTypeRepository(bindings.build(), typeRepository);
+        return EvaluationContext.forBindingsAndTypeRepository(bindings.build(), typeRepository, auxiliaryStores);
     }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/CascadesPlanner.java

@@ -61,6 +61,7 @@
 import com.apple.foundationdb.record.query.plan.cascades.matching.structure.PlannerBindings;
 import com.apple.foundationdb.record.query.plan.cascades.matching.structure.ReferenceMatchers;
 import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
+import com.apple.foundationdb.record.util.pair.NonnullPair;
 import com.google.common.base.Suppliers;
 import com.google.common.base.Verify;
 import com.google.common.collect.Iterables;
@@ -71,6 +72,7 @@
 import java.util.ArrayDeque;
 import java.util.Collection;
 import java.util.Deque;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -400,6 +402,38 @@ public QueryPlanResult planGraph(@Nonnull final Supplier<Reference> referenceSup
         }
     }
 
+    public QueryPlanResult planGraph(@Nonnull final Supplier<Reference> referenceSupplier,
+                                     @Nonnull final Optional<Collection<String>> allowedIndexesOptional,
+                                     @Nonnull final IndexQueryabilityFilter indexQueryabilityFilter,
+                                     @Nonnull final EvaluationContext evaluationContext,
+                                     @Nonnull final Map<SchemaIdentifier, NonnullPair<RecordMetaData, RecordStoreState>> additionalSchemas) {
+        try {
+            planPartial(referenceSupplier,
+                    rootReference ->
+                            MetaDataPlanContext.forRootReferenceWithAdditionalSchemas(configuration,
+                                    metaData,
+                                    recordStoreState,
+                                    matchCandidateRegistry,
+                                    rootReference,
+                                    allowedIndexesOptional,
+                                    indexQueryabilityFilter,
+                                    additionalSchemas
+                            ),
+                    evaluationContext);
+            final var plan = resultOrFail();
+            final var constraints = QueryPlanConstraint.collectConstraints(plan);
+            return new QueryPlanResult(plan,
+                    QueryPlanInfo.newBuilder()
+                            .put(QueryPlanInfoKeys.CONSTRAINTS, constraints)
+                            .put(QueryPlanInfoKeys.STATS_MAPS,
+                                    PlannerEventStatsCollector.flatMapCollector(PlannerEventStatsCollector::getStatsMaps)
+                                            .orElse(null))
+                            .build());
+        } finally {
+            PlannerEventListeners.dispatchOnDone();
+        }
+    }
+
     private RecordQueryPlan resultOrFail() {
         final Set<RelationalExpression> finalExpressions = currentRoot.getFinalExpressions();
         Verify.verify(finalExpressions.size() <= 1, "more than one variant present");

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/MetaDataPlanContext.java

@@ -31,6 +31,7 @@
 import com.apple.foundationdb.record.query.IndexQueryabilityFilter;
 import com.apple.foundationdb.record.query.RecordQuery;
 import com.apple.foundationdb.record.query.plan.RecordQueryPlannerConfiguration;
+import com.apple.foundationdb.record.util.pair.NonnullPair;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Lists;
 
@@ -39,6 +40,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -59,10 +61,20 @@ public class MetaDataPlanContext implements PlanContext {
     @Nonnull
     private final Set<MatchCandidate> matchCandidates;
 
+    @Nonnull
+    private final Map<MatchCandidate, SchemaIdentifier> matchCandidateSchemaMap;
+
     private MetaDataPlanContext(@Nonnull final RecordQueryPlannerConfiguration plannerConfiguration,
                                 @Nonnull final Set<MatchCandidate> matchCandidates) {
+        this(plannerConfiguration, matchCandidates, Map.of());
+    }
+
+    private MetaDataPlanContext(@Nonnull final RecordQueryPlannerConfiguration plannerConfiguration,
+                                @Nonnull final Set<MatchCandidate> matchCandidates,
+                                @Nonnull final Map<MatchCandidate, SchemaIdentifier> matchCandidateSchemaMap) {
         this.plannerConfiguration = plannerConfiguration;
         this.matchCandidates = ImmutableSet.copyOf(matchCandidates);
+        this.matchCandidateSchemaMap = Map.copyOf(matchCandidateSchemaMap);
     }
 
     @Nonnull
@@ -71,6 +83,12 @@ public RecordQueryPlannerConfiguration getPlannerConfiguration() {
         return plannerConfiguration;
     }
 
+    @Nonnull
+    @Override
+    public SchemaIdentifier getSchemaIdForMatchCandidate(@Nonnull final MatchCandidate candidate) {
+        return matchCandidateSchemaMap.getOrDefault(candidate, SchemaIdentifier.current());
+    }
+
     @Nullable
     private static KeyExpression commonPrimaryKey(@Nonnull Iterable<RecordType> recordTypes) {
         KeyExpression common = null;
@@ -179,6 +197,68 @@ public static PlanContext forRootReference(@Nonnull final RecordQueryPlannerConf
             return new MetaDataPlanContext(plannerConfiguration, ImmutableSet.of());
         }
 
+        return new MetaDataPlanContext(plannerConfiguration,
+                buildMatchCandidates(metaData, recordStoreState, matchCandidateRegistry,
+                        queriedRecordTypeNames, allowedIndexesOptional, indexQueryabilityFilter));
+    }
+
+    /**
+     * Build a plan context pooling match candidates from the primary schema and all additional schemas.
+     * Used when a query references tables from more than one schema.
+     *
+     * @param plannerConfiguration planner configuration
+     * @param metaData primary schema metadata
+     * @param recordStoreState primary schema store state
+     * @param matchCandidateRegistry registry for match candidates
+     * @param rootReference root reference of the query
+     * @param allowedIndexesOptional optional set of allowed index names
+     * @param indexQueryabilityFilter filter for queryable indexes
+     * @param additionalSchemas map from secondary schema identifier to (metadata, store state) pair
+     * @return a plan context with match candidates from all schemas
+     */
+    public static PlanContext forRootReferenceWithAdditionalSchemas(
+            @Nonnull final RecordQueryPlannerConfiguration plannerConfiguration,
+            @Nonnull final RecordMetaData metaData,
+            @Nonnull final RecordStoreState recordStoreState,
+            @Nonnull final IndexMatchCandidateRegistry matchCandidateRegistry,
+            @Nonnull final Reference rootReference,
+            @Nonnull final Optional<Collection<String>> allowedIndexesOptional,
+            @Nonnull final IndexQueryabilityFilter indexQueryabilityFilter,
+            @Nonnull final Map<SchemaIdentifier, NonnullPair<RecordMetaData, RecordStoreState>> additionalSchemas) {
+        final var queriedRecordTypeNames = recordTypes().evaluate(rootReference);
+        final ImmutableSet.Builder<MatchCandidate> allCandidates = ImmutableSet.builder();
+        final Map<MatchCandidate, SchemaIdentifier> schemaMap = new java.util.LinkedHashMap<>();
+
+        final var primaryTypeNames = queriedRecordTypeNames.stream()
+                .filter(name -> metaData.getRecordTypes().containsKey(name))
+                .collect(ImmutableSet.toImmutableSet());
+        if (!primaryTypeNames.isEmpty()) {
+            allCandidates.addAll(buildMatchCandidates(metaData, recordStoreState, matchCandidateRegistry,
+                    primaryTypeNames, allowedIndexesOptional, indexQueryabilityFilter));
+        }
+
+        for (final Map.Entry<SchemaIdentifier, NonnullPair<RecordMetaData, RecordStoreState>> entry : additionalSchemas.entrySet()) {
+            final SchemaIdentifier schemaId = entry.getKey();
+            final RecordMetaData secondaryMetaData = entry.getValue().getLeft();
+            final RecordStoreState secondaryState = entry.getValue().getRight();
+            final Set<String> allTypes = secondaryMetaData.getRecordTypes().keySet();
+            final ImmutableSet<MatchCandidate> secondaryCandidates = buildMatchCandidates(secondaryMetaData, secondaryState, matchCandidateRegistry,
+                    allTypes, allowedIndexesOptional, indexQueryabilityFilter);
+            allCandidates.addAll(secondaryCandidates);
+            secondaryCandidates.forEach(c -> schemaMap.put(c, schemaId));
+        }
+
+        return new MetaDataPlanContext(plannerConfiguration, allCandidates.build(), schemaMap);
+    }
+
+    @Nonnull
+    private static ImmutableSet<MatchCandidate> buildMatchCandidates(
+            @Nonnull final RecordMetaData metaData,
+            @Nonnull final RecordStoreState recordStoreState,
+            @Nonnull final IndexMatchCandidateRegistry matchCandidateRegistry,
+            @Nonnull final Set<String> queriedRecordTypeNames,
+            @Nonnull final Optional<Collection<String>> allowedIndexesOptional,
+            @Nonnull final IndexQueryabilityFilter indexQueryabilityFilter) {
         final var queriedRecordTypes =
                 queriedRecordTypeNames.stream().map(metaData::getRecordType).collect(Collectors.toList());
         final var indexList = Lists.<Index>newArrayList();
@@ -218,6 +298,6 @@ public static PlanContext forRootReference(@Nonnull final RecordQueryPlannerConf
                     .ifPresent(matchCandidatesBuilder::add);
         }
 
-        return new MetaDataPlanContext(plannerConfiguration, matchCandidatesBuilder.build());
+        return matchCandidatesBuilder.build();
     }
 }

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/cascades/PlanContext.java

@@ -53,6 +53,18 @@ public Set<MatchCandidate> getMatchCandidates() {
     @Nonnull
     Set<MatchCandidate> getMatchCandidates();
 
+    /**
+     * Returns the {@link SchemaIdentifier} for a match candidate, or {@link SchemaIdentifier#current()} if the
+     * candidate belongs to the primary (current) schema.
+     *
+     * @param candidate a match candidate
+     * @return the schema identifier for the candidate
+     */
+    @Nonnull
+    default SchemaIdentifier getSchemaIdForMatchCandidate(@Nonnull MatchCandidate candidate) {
+        return SchemaIdentifier.current();
+    }
+
     @Nonnull
     static PlanContext emptyContext() {
         return EMPTY_CONTEXT;