diff --git a/changelog/unreleased/aijoin-qparser.yml b/changelog/unreleased/aijoin-qparser.yml new file mode 100644 index 000000000000..c9b5c7f61842 --- /dev/null +++ b/changelog/unreleased/aijoin-qparser.yml @@ -0,0 +1,10 @@ +title: > + Introducing {!aijoin} query for query-time join with auxiliary index. +type: added +authors: + - name: Mikhail Khludnev + nick: mkhl +links: + - name: SOLR-18307 + url: https://issues.apache.org/jira/browse/SOLR-18307 + diff --git a/solr/core/src/java/org/apache/solr/search/join/AIJoinQParserPlugin.java b/solr/core/src/java/org/apache/solr/search/join/AIJoinQParserPlugin.java new file mode 100644 index 000000000000..dbe28bd1cc0e --- /dev/null +++ b/solr/core/src/java/org/apache/solr/search/join/AIJoinQParserPlugin.java @@ -0,0 +1,240 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.solr.search.join; + +import java.io.IOException; +import java.io.OutputStream; +import java.lang.invoke.MethodHandles; +import java.nio.file.Path; +import java.util.concurrent.ExecutorService; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.store.Directory; +import org.apache.solr.common.SolrException; +import org.apache.solr.common.params.CommonParams; +import org.apache.solr.common.params.SolrParams; +import org.apache.solr.common.util.NamedList; +import org.apache.solr.core.CloseHook; +import org.apache.solr.core.CoreContainer; +import org.apache.solr.core.DirectoryFactory.DirContext; +import org.apache.solr.core.SolrCore; +import org.apache.solr.request.SolrQueryRequest; +import org.apache.solr.request.SolrQueryRequestBase; +import org.apache.solr.request.SolrRequestInfo; +import org.apache.solr.response.QueryResponseWriter; +import org.apache.solr.response.SolrQueryResponse; +import org.apache.solr.search.QParser; +import org.apache.solr.search.QParserPlugin; +import org.apache.solr.search.SolrIndexSearcher; +import org.apache.solr.search.SyntaxError; +import org.apache.solr.search.join.aijoin.AIJoinIndex; +import org.apache.solr.util.RefCounted; +import org.apache.solr.util.plugin.SolrCoreAware; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Query parser exercising {@link AIJoinIndex} inside a {@link SolrCore}: it mimics {@link + * ScoreJoinQParserPlugin}'s local parameters, but resolves matches through the sidecar join index + * instead of {@link org.apache.lucene.search.join.JoinUtil}. Local parameters: + * + * + * + * Example: {@code q={!aijoin from=manu_id_s to=id fromIndex=products}foo}. + * + *

Unlike {@link ScoreJoinQParserPlugin.OtherCoreJoinQuery}, which only borrows the from-side + * searcher long enough to build a self-contained {@code Query} in {@code createWeight}, an {@link + * org.apache.solr.search.join.aijoin.AIJoinQuery} keeps reading the from-side searcher on every + * {@code scorerSupplier} call (it may lazily build missing pair columns per to-segment), so a + * cross-core from-searcher is pinned open for the whole request via {@link + * SolrRequestInfo#addCloseHook}, the same mechanism {@link + * org.apache.solr.search.JoinQuery.JoinQueryWeight} uses for the regular {@code {!join}}. + * + *

One {@link AIJoinIndex} is opened per core in {@link #inform(SolrCore)}, backed by a directory + * under the core's dataDir (configurable via the {@code dir} init parameter, resolved relative to + * dataDir unless absolute), and closed when the core closes. This sidecar always belongs to the + * "to" side core -- the one this plugin is registered in. + * + *

Why this implements {@link QueryResponseWriter}: {@link + * org.apache.solr.core.SolrResourceLoader}'s {@code awareCompatibility} allowlist (see SOLR-8311) + * only lets specific plugin base types implement {@link SolrCoreAware}, and {@code QParserPlugin} + * isn't one of them, so a plain {@code implements SolrCoreAware} fails core load with "Invalid + * 'Aware' object". {@code QueryResponseWriter} is on the allowlist and happens to be the cheapest + * interface there to satisfy (two abstract methods, both unreachable stubs below -- this class is + * never registered as a {@code }). This is safe here specifically because + * {@code QParserPlugin} instances are loaded once per core load/reload via {@link + * org.apache.solr.core.PluginBag}, exactly like the already-whitelisted {@link + * org.apache.solr.handler.component.SearchComponent} -- never created ad-hoc per request ({@link + * QParser#getParser(String, SolrQueryRequest)} resolves the already registered instance via {@code + * req.getCore().getQueryPlugin(name)}). + */ +public class AIJoinQParserPlugin extends QParserPlugin + implements QueryResponseWriter, SolrCoreAware { + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + /** + * Init parameter: directory holding the sidecar join index, resolved against the core's dataDir + * unless absolute. Defaults to {@value #DEFAULT_DIR}. + */ + public static final String DIR = "dir"; + + private static final String DEFAULT_DIR = "aijoin"; + + private String configuredDir = DEFAULT_DIR; + + private volatile AIJoinIndex joinIndex; + + @Override + public void init(NamedList args) { + super.init(args); + if (args != null && args.get(DIR) != null) { + configuredDir = args.get(DIR).toString(); + } + } + + @Override + public void inform(SolrCore core) { + Path path = Path.of(configuredDir); + if (!path.isAbsolute()) { + path = Path.of(core.getDataDir()).resolve(path); + } else { + core.getCoreContainer().assertPathAllowed(path); + } + Directory directory = null; + try { + directory = + core.getDirectoryFactory() + .get(path.toString(), DirContext.DEFAULT, core.getSolrConfig().indexConfig.lockType); + joinIndex = new AIJoinIndex(directory); + } catch (IOException | RuntimeException e) { + if (directory != null) { + try { + core.getDirectoryFactory().release(directory); + } catch (IOException releaseException) { + e.addSuppressed(releaseException); + } + } + throw new SolrException( + SolrException.ErrorCode.SERVER_ERROR, "Failed to open AIJoinIndex at " + path, e); + } + final Directory capturedDirectory = directory; + core.addCloseHook( + new CloseHook() { + @Override + public void preClose(SolrCore core) { + try { + joinIndex.close(); + } catch (IOException e) { + log.warn("Failed closing AIJoinIndex", e); + } finally { + try { + core.getDirectoryFactory().release(capturedDirectory); + } catch (IOException e) { + log.warn("Failed releasing AIJoinIndex directory {}", capturedDirectory, e); + } + } + } + }); + } + + // QueryResponseWriter stubs, unreachable: implemented only to satisfy SolrCoreAware's allowlist, + // see the class javadoc. This plugin is never registered as a . + + @Override + public void write( + OutputStream out, SolrQueryRequest request, SolrQueryResponse response, String contentType) { + throw new UnsupportedOperationException( + AIJoinQParserPlugin.class.getSimpleName() + + " is a QParserPlugin, not a QueryResponseWriter"); + } + + @Override + public String getContentType(SolrQueryRequest request, SolrQueryResponse response) { + throw new UnsupportedOperationException( + AIJoinQParserPlugin.class.getSimpleName() + + " is a QParserPlugin, not a QueryResponseWriter"); + } + + @Override + public QParser createParser( + String qstr, SolrParams localParams, SolrParams params, SolrQueryRequest req) { + return new QParser(qstr, localParams, params, req) { + @Override + public Query parse() throws SyntaxError { + if (joinIndex == null) { + throw new SolrException( + SolrException.ErrorCode.SERVER_ERROR, + "AIJoinQParserPlugin is not initialized; is it registered as a ?"); + } + final String fromField = getParam("from"); + final String toField = getParam("to"); + if (fromField == null || toField == null) { + throw new SyntaxError("aijoin query parser requires 'from' and 'to' local params"); + } + final String fromIndex = localParams.get("fromIndex"); + final String v = localParams.get(CommonParams.VALUE); + final String myCore = req.getCore().getCoreDescriptor().getName(); + + final Query fromQuery; + final IndexSearcher fromSearcher; + ExecutorService fromExecutor; + if (fromIndex != null && !fromIndex.equals(myCore)) { + CoreContainer container = req.getCoreContainer(); + String coreName = + ScoreJoinQParserPlugin.getCoreName( + fromIndex, container, req.getCore(), toField, fromField, localParams); + SolrCore fromCore = container.getCore(coreName); + if (fromCore == null) { + throw new SolrException( + SolrException.ErrorCode.BAD_REQUEST, "Cross-core join: no such core " + coreName); + } + SolrRequestInfo info = SolrRequestInfo.getRequestInfo(); + if (info == null) { + fromCore.close(); + throw new SolrException( + SolrException.ErrorCode.BAD_REQUEST, "Cross-core aijoin must have SolrRequestInfo"); + } + // released once this request completes: the from-side searcher is read on every + // scorerSupplier() call, not just while building this query, so it must outlive parse() + info.addCloseHook(fromCore); + try (SolrQueryRequestBase otherReq = new SolrQueryRequestBase(fromCore, params)) { + fromQuery = QParser.getParser(v, otherReq).getQuery(); + } + RefCounted fromRef = fromCore.getSearcher(false, true, null); + info.addCloseHook(fromRef::decref); + fromSearcher = fromRef.get(); + fromExecutor = (ExecutorService) fromCore.getCoreContainer().getIndexSearcherExecutor(); + } else { + fromQuery = subQuery(v, null).getQuery(); + fromSearcher = req.getSearcher(); + fromExecutor = (ExecutorService) req.getCoreContainer().getIndexSearcherExecutor(); + } + + return joinIndex.newJoinQuery(fromField, fromQuery, fromSearcher, toField, fromExecutor); + } + }; + } +} diff --git a/solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinDocWriter.java b/solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinDocWriter.java new file mode 100644 index 000000000000..fc803a88a418 --- /dev/null +++ b/solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinDocWriter.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.solr.search.join.aijoin; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.solr.search.join.aijoin.AIJoinUtil.JoinColumnModel; + +/** + * Sibling of {@code AIJoinColumnWriter} writing the same pair columns through the plain {@link + * Document} / {@link IndexWriter#addDocuments} API instead of {@code + * org.apache.lucene.document.column}. The whole batch is built as one in-memory {@link List} of + * {@code batchNumDocs} documents and handed to a single {@link IndexWriter#addDocuments} call: per + * its block semantics that list is indexed atomically, with no flush allowed to land in the middle + * of it, so -- exactly like {@code AIJoinColumnWriter}'s single {@code addBatch} -- the whole batch + * is guaranteed to end up doc-for-doc (list index == doc id) in one sidecar segment, keeping doc + * 0's edges and every from-doc id aligned the same way. + */ +final class AIJoinDocWriter extends AIJoinWriter { + + AIJoinDocWriter() {} + + @Override + void writeJoinColumns(IndexWriter writer, int batchNumDocs, Map mappings) + throws IOException { + List docs = new ArrayList<>(batchNumDocs); + for (int i = 0; i < batchNumDocs; i++) { + docs.add(new Document()); + } + for (Map.Entry entry : mappings.entrySet()) { + addJoinColumns(docs, entry.getValue(), entry.getKey()); + } + // a single block: IndexWriter guarantees no intermediate flush splits it across segments + writer.addDocuments(docs); + writer.commit(); + } + + /** + * Adds one pair's fields to {@code docs}: the doc-map field resolving from-side doc ids to + * to-side doc ids, spread across the batch's docs, and the edges companion fields, always added + * to doc 0 even when the pair maps nothing, so a once-built pair is detectable in the join index + * and never rebuilt. + */ + private static void addJoinColumns( // TODO don't write minusones columns for tombstones!! + List docs, JoinColumnModel mapping, String pairFieldName) throws IOException { + addOrdMap(docs, AIJoinUtil.TO_DOC_VAL_BY_FROM_DOCNUM + pairFieldName, mapping); + addEdges(docs, AIJoinUtil.FROM_EDGES_PREFIX + pairFieldName, mapping.edges().fromDocEdges()); + addEdges(docs, AIJoinUtil.TO_EDGES_PREFIX + pairFieldName, mapping.edges().toDocEdges()); + addEdges( + docs, AIJoinUtil.TO_COUNT_PREFIX + pairFieldName, new int[] {mapping.edges().toCount()}); + } + + /** + * Adds a pair's {min, max} (or count) values to doc 0, mirroring {@code AIJoinColumnWriter}'s + * {@code edgesColumn}, which puts both values at doc 0 too. + */ + private static void addEdges(List docs, String fieldName, int[] values) { + Document doc0 = docs.get(0); + for (int value : values) { + doc0.add(new SortedNumericDocValuesField(fieldName, value)); + } + } + + /** + * Adds the doc-map field: batch-local doc number is the from-side doc id and the SORTED_NUMERIC + * docvalue is the matching to-side doc id. From docs without a match get no value, hence the + * field is sparse. + */ + private static void addOrdMap(List docs, String fieldName, JoinColumnModel mapping) + throws IOException { + SortedNumericDocValues values = mapping.toDocByFromDoc(); + for (int doc = values.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = values.nextDoc()) { + docs.get(doc).add(new SortedNumericDocValuesField(fieldName, values.nextValue())); + } + } +} diff --git a/solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinIndex.java b/solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinIndex.java new file mode 100644 index 000000000000..63f6aeb3731a --- /dev/null +++ b/solr/core/src/java/org/apache/solr/search/join/aijoin/AIJoinIndex.java @@ -0,0 +1,447 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You 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 org.apache.solr.search.join.aijoin; + +import java.io.Closeable; +import java.io.IOException; +import java.lang.invoke.MethodHandles; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; +import org.apache.lucene.index.ConcurrentMergeScheduler; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.MergeScheduler; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.SearcherManager; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.CollectionUtil; +import org.apache.lucene.util.IOUtils; +import org.apache.solr.search.join.aijoin.AIJoinUtil.JoinColumnModel; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * The auxiliary join index: a self-maintaining sidecar persisting per (from-segment, to-segment) + * doc id mappings, so query-time joining reduces to bitset translation. It owns the sidecar's + * {@link IndexWriter} and {@link SearcherManager}; pair columns are built lazily when an {@link + * AIJoinQuery} first needs them, so users only construct an instance once, create queries with + * {@link #newJoinQuery} and search them with a bare to-side {@link IndexSearcher}: + * + * 

+ * AIJoinIndex joinIndex = new AIJoinIndex(joinDir);   // once per process
+ * Query q = joinIndex.newJoinQuery(fromField, fromQuery, fromSearcher, toField);
+ * TopDocs hits = toSearcher.search(q, 10);
+ * ...
+ * joinIndex.close();                                   // app shutdown
+ *
+ * + *

After either side reopens, the next query builds only the missing (from, to) segment pairs: + * pair columns are addressed by both sides' persistent segment keys, which survive reopens. Pair + * columns orphaned by merges are not reclaimed yet; see {@code README.md} in this package. + */ +public final class AIJoinIndex implements Closeable { + + private final IndexWriter writer; + private final SearcherManager manager; + + /** + * Dedups concurrent builders per pair field name: the thread that installs the future writes the + * pair, others wait on it. Completed futures stay put so a builder that raced a not-yet-visible + * refresh cannot write a duplicate pair column. + */ + private final ConcurrentHashMap>> + pairBuilds = new ConcurrentHashMap<>(); + + // package-private (not private): tests reach in directly to observe the reaper's state + final AIJoinMergePolicy mergePolicy; + private final MergeScheduler mergeScheduler; + static final AIJoinWriter INSTANCE = new AIJoinDocWriter(); // new AIJoinColumnWriter(); + + private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass()); + + /** Why a build was triggered; reported as {@code cause=} on the {@code AIJOIN evt=build} line. */ + enum BuildCause { + /** {nolink #ensureJoinSegments}, i.e. eagerly at {@link AIJoinQuery#createWeight} time. */ + EAGER_CREATE_WEIGHT, + /** {@link ToLeafJoinContext}, i.e. lazily for a gap the eager pass did not cover. */ + LAZY_TO_SEGMENT; + + @Override + public String toString() { + return name().toLowerCase(Locale.ROOT).replace('_', '-'); + } + } + + /** A pair's (from-segment, to-segment) leaf ordinals. */ + record SegmentsTuple(int fromLeafOrd, int toLeafOrd) {} + + /** + * A pair column's address in the join index: the pair field name and the sidecar segment (name + * plus current leaf ordinal) carrying it -- enough to locate and open the column's real + * docvalues, or to check whether a pair already exists before deciding what still needs to be + * built. A resolved cell's edges are tracked separately, as a plain {@code DocEdges}, since they + * don't change as this reference is refreshed. + */ + record JoinSegmentReference( + String pairFieldName, String joinSegmentName, int joinSegmentLeafOrd) {} + + /** + * Scans {@code joinSearcher}'s leaves for every pair column whose field name satisfies {@code + * isNeeded}, returning where each one lives. Used both to seed a fresh {@link AIJoinWeight}'s + * view of already-built pairs, and by {@link ToLeafJoinContext} to relocate a pair whose cached + * segment reference no longer resolves. TODO subject for in-heap caching TODO commit's userdata + * might have a list of pairs with known segment ords and names + */ + static Map extractExistingJoinColumns( + IndexSearcher joinSearcher, Predicate isNeeded) { + Map existingJoinSegments = + CollectionUtil.newHashMap(joinSearcher.getIndexReader().leaves().size()); + for (LeafReaderContext joinContext : joinSearcher.getIndexReader().leaves()) { + String segmentName = AIJoinUtil.segmentName(joinContext); + for (FieldInfo fieldInfo : joinContext.reader().getFieldInfos()) { + // pairs are detected by their toCount companion, which is written to doc 0 for every + // built pair; the join column itself is sparse and a tombstone pair (disjoint terms) + // never materializes it, so scanning for join columns kept re-reporting once-built + // tombstones as missing -- and re-triggering their from-side FK loads on every query + String splits[] = fieldInfo.name.split(AIJoinUtil.TO_COUNT_PREFIX); + if (splits.length == 2 && isNeeded.test(splits[1])) { + existingJoinSegments.computeIfAbsent( + splits[1], + fieldName -> new JoinSegmentReference(fieldName, segmentName, joinContext.ord)); + } + } + } + return existingJoinSegments; + } + + /** + * Opens a persistent auxiliary join index over the given directory, creating it if empty, using + * the default {@link AIJoinIndexConfig}. The caller retains ownership of the directory: {@link + * #close()} does not close it. + */ + public AIJoinIndex(Directory directory) throws IOException { + this(directory, new AIJoinIndexConfig()); + } + + /** + * Opens a persistent auxiliary join index over the given directory, creating it if empty, using + * the given {@link AIJoinIndexConfig}. The caller retains ownership of the directory: {@link + * #close()} does not close it. + */ + public AIJoinIndex(Directory directory, AIJoinIndexConfig config) throws IOException { + this(directory, config, new ConcurrentMergeScheduler()); + } + + /** + * Opens a persistent auxiliary join index over the given directory, creating it if empty, using + * the given {@link AIJoinIndexConfig} and {@link MergeScheduler} in place of the default {@link + * ConcurrentMergeScheduler}. The caller retains ownership of the directory: {@link #close()} does + * not close it. + */ + public AIJoinIndex(Directory directory, AIJoinIndexConfig config, MergeScheduler mergeScheduler) + throws IOException { + this.mergeScheduler = mergeScheduler; + this.mergePolicy = new AIJoinMergePolicy(); + this.mergePolicy.setSweepInterval(config.getSweepSamplingIntervalNanos(), TimeUnit.NANOSECONDS); + this.writer = + new IndexWriter( + directory, + new IndexWriterConfig().setMergePolicy(mergePolicy).setMergeScheduler(mergeScheduler)); + this.manager = new SearcherManager(writer, null); + } + + /** + * Creates a query joining the docs matching {@code fromQuery} in {@code fromSearcher}'s index to + * the index the returned query is executed against, through {@code fromField} = {@code toField} + * term equality. Missing pair columns are built into this join index on first execution. + * + * @deprecated use another constructor passing executor service + */ + @Deprecated + public Query newJoinQuery( + String fromField, Query fromQuery, IndexSearcher fromSearcher, String toField) { + return newJoinQuery(fromField, fromQuery, fromSearcher, toField, new DirectExecutorService()); + } + + public Query newJoinQuery( + String fromField, + Query fromQuery, + IndexSearcher fromSearcher, + String toField, + ExecutorService fromExecutor) { + return new AIJoinQuery( + this, + fromField, + fromQuery, + fromSearcher, + toField, + fromExecutor == null ? new DirectExecutorService() : fromExecutor); + } + + /** + * How many of the given pair field names already have a build claimed (in-flight or completed) in + * {@link #pairBuilds}. Diagnostic only: a pair counted here but still reported missing by {@link + * #extractExistingJoinColumns} means the caller is about to redo from-side work for a pair that + * was already built -- its column just isn't visible through the searcher it consulted. + */ + int countClaimedBuilds(Set pairFieldNames) { + int claimed = 0; + for (String pairFieldName : pairFieldNames) { + if (pairBuilds.containsKey(pairFieldName)) { + claimed++; + } + } + return claimed; + } + + IndexSearcher acquire() throws IOException { + return manager.acquire(); + } + + void release(IndexSearcher searcher) throws IOException { + manager.release(searcher); + } + + /** + * Builds and persists the given missing pair columns, keyed by pair field name to their + * (from-segment, to-segment) leaf ordinals. Pairs concurrently built by another thread are + * awaited, not rebuilt. On return the internal searcher manager is refreshed past every requested + * pair. + * + * @return in memory data for just written segemts + */ + Map writeJoinSegments( + Map missingPairs, + IndexReader fromReader, + String fromField, + IndexReader toReader, + String toField, + BuildCause buildCause, + String ctxId, + Future[] fromColumnFutures) + throws IOException, ExecutionException, InterruptedException { + long startNanos = System.nanoTime(); + int batchNumDocsLogged = 0; + Map>> owned = + new LinkedHashMap<>(); + List>> awaited = new ArrayList<>(); + for (String pairFieldName : missingPairs.keySet()) { + CompletableFuture> created = new CompletableFuture<>(); + CompletableFuture> existing = + pairBuilds.putIfAbsent(pairFieldName, created); + if (existing == null) { + owned.put(pairFieldName, created); + } else { + awaited.add(existing); + } + } + Map loadedMappings = new LinkedHashMap<>(); + try { + if (!owned.isEmpty()) { + // all owned pairs go into a single batch: pair columns are addressed by from-side doc id, + // so a batch must start at doc 0 of its sidecar segment, which writeBatch guarantees by + // flushing one batch per commit + int batchNumDocs = 0; + for (String pairFieldName : owned.keySet()) { + SegmentsTuple position = missingPairs.get(pairFieldName); + LeafReaderContext toContext = toReader.leaves().get(position.toLeafOrd()); + LeafReaderContext fromContext = fromReader.leaves().get(position.fromLeafOrd()); + assert fromColumnFutures[fromContext.ord] != null; + AIJoinUtil.JoinColumnModel mapping = + AIJoinUtil.computeDocMapping( + toContext, + toField, // new ForeignKeyColumn(fromContext, fromField) + fromColumnFutures[fromContext.ord].get().fkColumn); + batchNumDocs = Math.max(batchNumDocs, fromContext.reader().maxDoc()); + loadedMappings.put(pairFieldName, mapping); + } + writeBatch(batchNumDocs, loadedMappings); + batchNumDocsLogged = batchNumDocs; + // TODO flush every single field to get single field segments + for (Map.Entry>> entry : + owned.entrySet()) { + entry + .getValue() + .complete( + new AbstractMap.SimpleImmutableEntry<>( + entry.getKey(), loadedMappings.get(entry.getKey()))); + } + } + } catch (Throwable t) { + // withdraw the claims so a later query can retry the build + for (Map.Entry>> entry : + owned.entrySet()) { + pairBuilds.remove(entry.getKey(), entry.getValue()); + entry.getValue().completeExceptionally(t); + } + throw t; + } + long builtNanos = System.nanoTime() - startNanos; + Map result = new LinkedHashMap<>(loadedMappings); + for (CompletableFuture> future : awaited) { + try { + Map.Entry entry = future.join(); + result.put(entry.getKey(), entry.getValue()); + } catch (CompletionException e) { + Throwable cause = e.getCause(); + if (cause instanceof IOException ioe) { + throw ioe; + } + if (cause instanceof RuntimeException re) { + throw re; + } + throw new IOException(cause); + } + } + if (AIJoinUtil.diagnosticsEnabled(log) && !missingPairs.isEmpty()) { + long toCount = 0; + for (JoinColumnModel model : loadedMappings.values()) { + toCount += model.edges().toCount(); + } + // built/awaited split matters: an awaited pair cost this thread only the wait, so folding + // the two together would attribute another thread's build work to this query + AIJoinUtil.logDiagnostic( + log, + "AIJOIN evt=build ctx={} cause={} pairsRequested={} pairsBuilt={} pairsAwaited={}" + + " builtMs={} awaitedMs={} toCount={} batchNumDocs={} writtenPairs={}", + ctxId == null ? "-" : ctxId, + buildCause, + missingPairs.size(), + loadedMappings.size(), + awaited.size(), + builtNanos / 1_000_000L, + (System.nanoTime() - startNanos - builtNanos) / 1_000_000L, + toCount, + batchNumDocsLogged, + loadedMappings.keySet()); + } + return result; + } + + /** + * deprecated don't write all of them upfront + * + *

Eagerly builds and persists every pair column in {@code neededPairs} not yet present in this + * join index, so an {@link AIJoinWeight} being constructed at {@link AIJoinQuery#createWeight} + * already sees a complete view of the pairs it needs, instead of discovering gaps lazily -- one + * to-segment at a time -- in {@link ToLeafJoinContext}. Missing pairs are resolved to their + * (from-segment, to-segment) leaf ordinals by crossing {@code fromSearcher}'s leaves against + * {@code toSearcher}'s leaves; pairs concurrently built by another thread are awaited, not + * rebuilt (see {@link #writeJoinSegments}). + */ + // @Deprecated + // void ensureJoinSegments( + // Set neededPairs, + // IndexSearcher fromSearcher, + // String fromField, + // IndexSearcher toSearcher, + // String toField) + // throws IOException { + // Map existing; + // IndexSearcher joinSearcher = acquire(); + // try { + // existing = extractExistingJoinColumns(joinSearcher, neededPairs::contains); + // } finally { + // release(joinSearcher); + // } + // if (existing.keySet().containsAll(neededPairs)) { + // return; + // } + // Map missingPairs = new LinkedHashMap<>(); + // for (LeafReaderContext fromContext : fromSearcher.getLeafContexts()) { + // for (LeafReaderContext toContext : toSearcher.getLeafContexts()) { + // String pairFieldName = AIJoinUtil.pairFieldName(fromContext, fromField, toContext, + // toField); + // if (neededPairs.contains(pairFieldName) && !existing.containsKey(pairFieldName)) { + // missingPairs.put(pairFieldName, new SegmentsTuple(fromContext.ord, toContext.ord)); + // } + // } + // } + // if (!missingPairs.isEmpty()) { + // writeJoinSegments( + // missingPairs, + // fromSearcher.getIndexReader(), + // fromField, + // toSearcher.getIndexReader(), + // toField, + // BuildCause.EAGER_CREATE_WEIGHT, + // null, fromColumnFutures); // runs before any ToLeafJoinContext exists, so there is no + // context to blame + // } + // } + + /** + * Serializes sidecar writes: one batch per commit keeps every batch at doc 0 of its own segment, + * preserving pair-column doc number == from-side doc id. Completing builders' futures after this + * returns guarantees waiters observe the refreshed reader. + * + *

It should be plain simple synchronized. As alternatives + * + *