auxjoin qparser - a prunable segment-parallel auxiliary index join

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
+

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:
+ *
+ * <ul>
+ *   <li>from - "foreign key" field name, collected while enumerating the subordinate query (the
+ *       local parameter value).
+ *   <li>fromIndex - optional core name to run the subordinate query against, when it differs from
+ *       this core; cross-core joins are the reason {@link AIJoinIndex} exists in the first place,
+ *       so this mirrors {@link ScoreJoinQParserPlugin}'s <code>fromIndex</code>, including
+ *       SolrCloud alias/collection resolution via {@link ScoreJoinQParserPlugin#getCoreName}.
+ *   <li>to - "primary key" field name looked up in this core's index.
+ * </ul>
+ *
+ * Example: {@code q={!aijoin from=manu_id_s to=id fromIndex=products}foo}.
+ *
+ * <p>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}}.
+ *
+ * <p>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.
+ *
+ * <p><b>Why this implements {@link QueryResponseWriter}:</b> {@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 <queryResponseWriter>}). 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 <queryResponseWriter>.
+
+  @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 <queryParser>?");
+        }
+        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<SolrIndexSearcher> 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);
+      }
+    };
+  }
+}

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<String, JoinColumnModel> mappings)
+      throws IOException {
+    List<Document> docs = new ArrayList<>(batchNumDocs);
+    for (int i = 0; i < batchNumDocs; i++) {
+      docs.add(new Document());
+    }
+    for (Map.Entry<String, JoinColumnModel> 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<Document> 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<Document> 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<Document> 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()));
+    }
+  }
+}