Merge pull request #30717 from vespa-engine/balder/cluster.schema-map-to.cluster

Balder/cluster.schema map to.cluster

6 files changed, 91 insertions, 51 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java b/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
index 000fef9b14b..72184c5ea32 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/FederationSearcher.java
@@ -18,6 +18,7 @@ import com.yahoo.search.Result;
 import com.yahoo.search.Searcher;
 import com.yahoo.search.federation.selection.FederationTarget;
 import com.yahoo.search.federation.selection.TargetSelector;
+import com.yahoo.search.federation.sourceref.ModifyQueryAndResult;
 import com.yahoo.search.federation.sourceref.SearchChainInvocationSpec;
 import com.yahoo.search.federation.sourceref.SearchChainResolver;
 import com.yahoo.search.federation.sourceref.SingleTarget;
@@ -54,6 +55,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static com.yahoo.collections.CollectionUtil.first;
@@ -88,7 +90,7 @@ public class FederationSearcher extends ForkingSearcher {
     public FederationSearcher(FederationConfig config, SchemaInfo schemaInfo,
                               ComponentRegistry<TargetSelector> targetSelectors) {
         this(createResolver(config),
-             VirtualSourceResolver.of(config),
+             createVirtualSourceResolver(config),
              resolveSelector(config.targetSelector(), targetSelectors),
              createSchema2Clusters(schemaInfo));
     }
@@ -109,6 +111,10 @@ public class FederationSearcher extends ForkingSearcher {
         this.virtualSourceResolver = virtualSourceResolver;
     }
 
+    private static VirtualSourceResolver createVirtualSourceResolver(FederationConfig config) {
+        return VirtualSourceResolver.of(config.target().stream().map(FederationConfig.Target::id).collect(Collectors.toUnmodifiableSet()));
+    }
+
     private static TargetSelector resolveSelector(String selectorId,
                                                   ComponentRegistry<TargetSelector> targetSelectors) {
         if (selectorId.isEmpty()) return null;
@@ -149,15 +155,36 @@ public class FederationSearcher extends ForkingSearcher {
         return builder.build();
     }
 
+    private static class SearchChaininvocationProxy extends SearchChainInvocationSpec {
+        SearchChaininvocationProxy(ComponentId searchChainId, FederationOptions federationOptions, String schema) {
+            super(searchChainId, federationOptions, List.of(schema));
+        }
+
+        @Override
+        public void modifyTargetQuery(Query query) {
+            query.getModel().setSources(searchChainId.getName());
+            query.getModel().setRestrict(schemas.get(0));
+        }
+    }
+
     private static void addSearchChain(SearchChainResolver.Builder builder,
                                        FederationConfig.Target target,
                                        FederationConfig.Target.SearchChain searchChain)
     {
-        if (!target.id().equals(searchChain.searchChainId()))
-            throw new RuntimeException("Invalid federation config, " + target.id() + " != " + searchChain.searchChainId());
-
-        builder.addSearchChain(ComponentId.fromString(searchChain.searchChainId()),
-                               federationOptions(searchChain), searchChain.documentTypes());
+        String id = target.id();
+        if (!id.equals(searchChain.searchChainId()))
+            throw new RuntimeException("Invalid federation config, " + id + " != " + searchChain.searchChainId());
+
+        ComponentId searchChainId = ComponentId.fromString(id);
+        builder.addSearchChain(searchChainId, federationOptions(searchChain), searchChain.documentTypes());
+        // Here we make synthetic SearchChain proxies for all cluster.schema combinations possible
+        // Given a source on the form saerchcluster.schema will rewrite it to source=searchcluster and restrict to schema.
+        // TODO Consider solving this in the config model by making many synthetic search clusters
+        for (String schema : searchChain.documentTypes()) {
+            String virtualChainId = id + "." + schema;
+            builder.addSearchChain(ComponentId.fromString(virtualChainId),
+                    new SearchChaininvocationProxy(searchChainId, federationOptions(searchChain).setUseByDefault(false), schema));
+        }
     }
 
     private static void addSourceForProvider(SearchChainResolver.Builder builder, FederationConfig.Target target,
@@ -209,7 +236,6 @@ public class FederationSearcher extends ForkingSearcher {
 
     private void search(Query query, Execution execution, Target target, Result mergedResults) {
         mergeResult(query, target, mergedResults, search(query, execution, target).orElse(createSearchChainTimedOutResult(query, target)));
-
     }
 
     private void search(Query query, Execution execution, Collection<Target> targets, Result mergedResults) {
@@ -577,7 +603,7 @@ public class FederationSearcher extends ForkingSearcher {
         Collection<SearchChainInvocationSpec> prunedTargets = new ArrayList<>();
 
         for (SearchChainInvocationSpec target : targets) {
-            if (target.documentTypes.isEmpty() || documentTypeIntersectionIsNonEmpty(restrict, target))
+            if (target.schemas.isEmpty() || documentTypeIntersectionIsNonEmpty(restrict, target))
                 prunedTargets.add(target);
         }
 
@@ -585,7 +611,7 @@ public class FederationSearcher extends ForkingSearcher {
     }
 
     private boolean documentTypeIntersectionIsNonEmpty(Set<String> restrict, SearchChainInvocationSpec target) {
-        for (String documentType : target.documentTypes) {
+        for (String documentType : target.schemas) {
             if (restrict.contains(documentType))
                 return true;
         }
@@ -615,11 +641,9 @@ public class FederationSearcher extends ForkingSearcher {
     }
 
     /** A target for federation, containing a chain to which a federation query can be forwarded. */
-    static abstract class Target {
+    static abstract class Target implements ModifyQueryAndResult {
 
         abstract Chain<Searcher> getChain();
-        abstract void modifyTargetQuery(Query query);
-        abstract void modifyTargetResult(Result result);
 
         ComponentId getId() {
             return getChain().getId();
@@ -650,9 +674,9 @@ public class FederationSearcher extends ForkingSearcher {
         Chain<Searcher> getChain() { return chain; }
 
         @Override
-        void modifyTargetQuery(Query query) {}
+        public void modifyTargetQuery(Query query) { target.modifyTargetQuery(query); }
         @Override
-        void modifyTargetResult(Result result) {}
+        public void modifyTargetResult(Result result) { target.modifyTargetResult(result); }
 
         @Override
         public FederationOptions federationOptions() { return target.federationOptions; }
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/ModifyQueryAndResult.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/ModifyQueryAndResult.java
new file mode 100644
index 00000000000..7a9cf1b561f
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/ModifyQueryAndResult.java
@@ -0,0 +1,9 @@
+package com.yahoo.search.federation.sourceref;
+
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
+
+public interface ModifyQueryAndResult {
+    void modifyTargetQuery(Query query);
+    void modifyTargetResult(Result result);
+}
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java
index 32af08f95fc..9822c4ea09a 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainInvocationSpec.java
@@ -2,6 +2,8 @@
 package com.yahoo.search.federation.sourceref;
 
 import com.yahoo.component.ComponentId;
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
 import com.yahoo.search.searchchain.model.federation.FederationOptions;
 
 import java.util.List;
@@ -13,7 +15,7 @@ import java.util.Objects;
  *
  * @author Tony Vaagenes
  */
-public class SearchChainInvocationSpec implements Cloneable {
+public class SearchChainInvocationSpec implements ModifyQueryAndResult, Cloneable {
 
     public final ComponentId searchChainId;
 
@@ -23,16 +25,20 @@ public class SearchChainInvocationSpec implements Cloneable {
     public final ComponentId provider;
 
     public final FederationOptions federationOptions;
-    public final List<String> documentTypes;
+    public final List<String> schemas;
 
-    SearchChainInvocationSpec(ComponentId searchChainId,
-                              ComponentId source, ComponentId provider, FederationOptions federationOptions,
-                              List<String> documentTypes) {
+    public SearchChainInvocationSpec(ComponentId searchChainId, FederationOptions federationOptions, List<String> schemas) {
+        this(searchChainId, null, null, federationOptions, schemas);
+    }
+
+
+    SearchChainInvocationSpec(ComponentId searchChainId, ComponentId source, ComponentId provider,
+                              FederationOptions federationOptions, List<String> schemas) {
         this.searchChainId = searchChainId;
         this.source = source;
         this.provider = provider;
         this.federationOptions = federationOptions;
-        this.documentTypes = List.copyOf(documentTypes);
+        this.schemas = List.copyOf(schemas);
     }
 
     @Override
@@ -49,13 +55,16 @@ public class SearchChainInvocationSpec implements Cloneable {
         if ( ! Objects.equals(this.source, other.source)) return false;
         if ( ! Objects.equals(this.provider, other.provider)) return false;
         if ( ! Objects.equals(this.federationOptions, other.federationOptions)) return false;
-        if ( ! Objects.equals(this.documentTypes, other.documentTypes)) return false;
+        if ( ! Objects.equals(this.schemas, other.schemas)) return false;
         return true;
     }
 
     @Override
     public int hashCode() { 
-        return Objects.hash(searchChainId, source, provider, federationOptions, documentTypes);
+        return Objects.hash(searchChainId, source, provider, federationOptions, schemas);
     }
 
+    public void modifyTargetQuery(Query query) { }
+    public void modifyTargetResult(Result result) {}
+
 }
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java
index df91b968750..7dc65c819e4 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/SearchChainResolver.java
@@ -46,6 +46,16 @@ public class SearchChainResolver {
 
     public static class Builder {
 
+        public interface InvocationSpecFactory {
+            SearchChainInvocationSpec create(ComponentId searchChainId, FederationOptions federationOptions, List<String> schemas);
+        }
+
+        private class DefaultInvocationSpecFactory implements InvocationSpecFactory {
+            public SearchChainInvocationSpec create(ComponentId searchChainId, FederationOptions federationOptions, List<String> schemas) {
+                return new SearchChainInvocationSpec(searchChainId, federationOptions, schemas);
+            }
+        }
+
         private final SortedSet<Target> defaultTargets = new TreeSet<>();
 
         private final ComponentRegistry<Target> targets = new ComponentRegistry<>() {
@@ -60,22 +70,22 @@ public class SearchChainResolver {
             return addSearchChain(searchChainId, federationOptions, List.of());
         }
 
-        public Builder addSearchChain(ComponentId searchChainId, List<String> documentTypes) {
-            return addSearchChain(searchChainId, new FederationOptions(), documentTypes);
+        public Builder addSearchChain(ComponentId searchChainId, List<String> schemas) {
+            return addSearchChain(searchChainId, new FederationOptions(), schemas);
         }
 
         public Builder addSearchChain(ComponentId searchChainId,
                                       FederationOptions federationOptions,
-                                      List<String> documentTypes) {
-            registerTarget(new SingleTarget(searchChainId,
-                                            new SearchChainInvocationSpec(searchChainId,
-                                                                          null,
-                                                                          null,
-                                                                          federationOptions,
-                                                                          documentTypes),
-                                            false));
+                                      List<String> schemas) {
+            addSearchChain(new SearchChainInvocationSpec(searchChainId, federationOptions, schemas));
             return this;
         }
+        private Builder addSearchChain(SearchChainInvocationSpec invocationSpec) {
+            return registerTarget(new SingleTarget(invocationSpec.searchChainId, invocationSpec, false));
+        }
+        public Builder addSearchChain(ComponentId id, SearchChainInvocationSpec invocationSpec) {
+            return registerTarget(new SingleTarget(id, invocationSpec, false));
+        }
 
         private Builder registerTarget(SingleTarget singleTarget) {
             targets.register(singleTarget.getId(), singleTarget);
@@ -87,9 +97,8 @@ public class SearchChainResolver {
 
         public Builder addSourceForProvider(ComponentId sourceId, ComponentId providerId, ComponentId searchChainId,
                                             boolean isDefaultProviderForSource, FederationOptions federationOptions,
-                                            List<String> documentTypes) {
-            SearchChainInvocationSpec searchChainInvocationSpec =
-                    new SearchChainInvocationSpec(searchChainId, sourceId, providerId, federationOptions, documentTypes);
+                                            List<String> schemas) {
+            var searchChainInvocationSpec = new SearchChainInvocationSpec(searchChainId, sourceId, providerId, federationOptions, schemas);
 
             SourcesTarget sourcesTarget = getOrRegisterSourceTarget(sourceId);
             sourcesTarget.addSource(providerId, searchChainInvocationSpec, isDefaultProviderForSource);
diff --git a/container-search/src/main/java/com/yahoo/search/federation/sourceref/VirtualSourceResolver.java b/container-search/src/main/java/com/yahoo/search/federation/sourceref/VirtualSourceResolver.java
index fc07d12d429..2b2c3b2f3ca 100644
--- a/container-search/src/main/java/com/yahoo/search/federation/sourceref/VirtualSourceResolver.java
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/VirtualSourceResolver.java
@@ -1,8 +1,6 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.federation.sourceref;
 
-import com.yahoo.search.federation.FederationConfig;
-
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
@@ -43,9 +41,7 @@ public class VirtualSourceResolver {
         }
         return virtualSourceMap;
     }
-    public static VirtualSourceResolver of(FederationConfig config) {
-        return of(config.target().stream().map(FederationConfig.Target::id).collect(Collectors.toUnmodifiableSet()));
-    }
+
     public Set<String> resolve(Set<String> sourcesInQuery) {
         boolean hasMapping = sourcesInQuery.stream().anyMatch(virtualSources::containsKey);
         if (hasMapping) {
diff --git a/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java b/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java
index bd39aed38fa..b1ce1fbd3ee 100644
--- a/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java
+++ b/container-search/src/test/java/com/yahoo/search/federation/FederationResultTest.java
@@ -120,17 +120,10 @@ public class FederationResultTest {
             this.timeout = timeout;
         }
 
-        @Override
-        Chain<Searcher> getChain() { return chain; }
-
-        @Override
-        void modifyTargetQuery(Query query) { }
-
-        @Override
-        void modifyTargetResult(Result result) { }
-
-        @Override
-        public FederationOptions federationOptions() {
+        @Override Chain<Searcher> getChain() { return chain; }
+        @Override public void modifyTargetQuery(Query query) { }
+        @Override public void modifyTargetResult(Result result) { }
+        @Override public FederationOptions federationOptions() {
             return new FederationOptions(false, timeout, true);
         }