Merge pull request #30272 from vespa-engine/balder/resolve-virtual-sources

- Add a resolver for virtual sources.

3 files changed, 77 insertions, 10 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 e92fd692cac..1f8f8757ebc 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
@@ -25,6 +25,7 @@ import com.yahoo.search.federation.sourceref.SingleTarget;
 import com.yahoo.search.federation.sourceref.SourceRefResolver;
 import com.yahoo.search.federation.sourceref.SourcesTarget;
 import com.yahoo.search.federation.sourceref.UnresolvedSearchChainException;
+import com.yahoo.search.federation.sourceref.VirtualSourceResolver;
 import com.yahoo.search.query.Properties;
 import com.yahoo.search.result.ErrorMessage;
 import com.yahoo.search.result.Hit;
@@ -77,25 +78,28 @@ public class FederationSearcher extends ForkingSearcher {
 
     private final SearchChainResolver searchChainResolver;
     private final SourceRefResolver sourceRefResolver;
+    private final VirtualSourceResolver virtualSourceResolver;
 
     private final TargetSelector<?> targetSelector;
     private final Clock clock = Clock.systemUTC();
 
     @Inject
     public FederationSearcher(FederationConfig config, ComponentRegistry<TargetSelector> targetSelectors) {
-        this(createResolver(config), resolveSelector(config.targetSelector(), targetSelectors));
+        this(createResolver(config), VirtualSourceResolver.of(config), resolveSelector(config.targetSelector(), targetSelectors));
     }
 
     // for testing
     public FederationSearcher(ComponentId id, SearchChainResolver searchChainResolver) {
-        this(searchChainResolver, null);
+        this(searchChainResolver, VirtualSourceResolver.of(), null);
     }
 
     private FederationSearcher(SearchChainResolver searchChainResolver,
+                               VirtualSourceResolver virtualSourceResolver,
                                TargetSelector targetSelector) {
         this.searchChainResolver = searchChainResolver;
         sourceRefResolver = new SourceRefResolver(searchChainResolver);
         this.targetSelector = targetSelector;
+        this.virtualSourceResolver = virtualSourceResolver;
     }
 
     private static TargetSelector resolveSelector(String selectorId,
@@ -370,9 +374,9 @@ public class FederationSearcher extends ForkingSearcher {
             else { // fill timed out: Remove these hits as they are incomplete and may cause a race when accessed later
                 result.hits().addError(futureFilledResult.getSecond().createTimeoutError());
                 for (Iterator<Hit> i = futureFilledResult.getFirst().hits().unorderedDeepIterator(); i.hasNext(); ) {
-                    // Note that some of these hits may be filled, but as the fill thread may still be working on them
-                    // and we do not synchronize with it we need to discard all
-                    Hit removed = result.hits().remove(i.next().getId());
+                    // Note that some of these hits may be filled, but as the fill thread may still be working on them,
+                    // and we do not synchronize with it, we need to discard all.
+                    result.hits().remove(i.next().getId());
                 }
             }
         }
@@ -427,8 +431,10 @@ public class FederationSearcher extends ForkingSearcher {
                 resolveSources(sources, properties, indexFacts);
     }
 
-    private Results<SearchChainInvocationSpec, UnresolvedSearchChainException> resolveSources(Set<String> sources, Properties properties, IndexFacts indexFacts) {
+
+    private Results<SearchChainInvocationSpec, UnresolvedSearchChainException> resolveSources(Set<String> sourcesInQuery, Properties properties, IndexFacts indexFacts) {
         Results.Builder<SearchChainInvocationSpec, UnresolvedSearchChainException> result = new Builder<>();
+        Set<String> sources = virtualSourceResolver.resolve(sourcesInQuery);
 
         for (String source : sources) {
             try {
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 a2bc12ddbd0..df91b968750 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
@@ -56,10 +56,6 @@ public class SearchChainResolver {
             }
         };
 
-        public Builder addSearchChain(ComponentId searchChainId) {
-            return addSearchChain(searchChainId, List.of());
-        }
-
         public Builder addSearchChain(ComponentId searchChainId, FederationOptions federationOptions) {
             return addSearchChain(searchChainId, federationOptions, List.of());
         }
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
new file mode 100644
index 00000000000..fc07d12d429
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/federation/sourceref/VirtualSourceResolver.java
@@ -0,0 +1,65 @@
+// 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;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Multiple sources like contentcluster.schema1, contencluster.schema2 needs to respond
+ * when source is the prefix contentcluster. This is done by generating map from virtual source
+ * to the fully qualified ones, and resolving from there.
+ *
+ * @author baldersheim
+ */
+public class VirtualSourceResolver {
+    private final Map<String, Set<String>> virtualSources;
+    private VirtualSourceResolver(Map<String, Set<String>> virtualSources) {
+        this.virtualSources = virtualSources;
+    }
+    public static VirtualSourceResolver of() {
+        return new VirtualSourceResolver(Map.of());
+    }
+    public static VirtualSourceResolver of(Set<String> targets) {
+        return new VirtualSourceResolver(createVirtualSources(targets));
+    }
+    private static Map<String, Set<String>> createVirtualSources(Set<String> targets) {
+        Set<String> virtualSources = targets.stream()
+                .filter(id -> id.contains("."))
+                .map(id -> id.substring(0, id.indexOf('.')))
+                .collect(Collectors.toUnmodifiableSet());
+        if (virtualSources.isEmpty()) return Map.of();
+        Map<String, Set<String>> virtualSourceMap = new HashMap<>();
+        for (String virtualSource : virtualSources) {
+            String prefix = virtualSource + ".";
+            Set<String> sources = targets.stream()
+                    .filter(id -> id.startsWith(prefix))
+                    .collect(Collectors.toUnmodifiableSet());
+            virtualSourceMap.put(virtualSource, sources);
+        }
+        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) {
+            Set<String> resolved = new HashSet<>();
+            for (String source : sourcesInQuery) {
+                var subSources = virtualSources.get(source);
+                if (subSources != null) {
+                    resolved.addAll(subSources);
+                } else {
+                    resolved.add(source);
+                }
+            }
+            return resolved;
+        }
+        return sourcesInQuery;
+    }
+}