Promote PerThreadPKLookup to core as PrimaryKeyLookup

lucene/CHANGES.txt

@@ -119,6 +119,9 @@ New Features
 
 * GITHUB#15818: Add BM25 k3 query-term frequency saturation to BM25Similarity. (Sagar Upadhyaya)
 
+* GITHUB#15520: Promote PerThreadPKLookup from the test-framework to core as
+  PrimaryKeyLookup. The original PerThreadPKLookup is now deprecated. (Shubham Sharma)
+
 Improvements
 ---------------------
 * GITHUB#15704: Replace LinkedList with more efficient data structure. (Renato Haeberli)

lucene/core/src/java/org/apache/lucene/index/PrimaryKeyLookup.java

@@ -0,0 +1,185 @@
+/*
+ * 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.lucene.index;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.lucene.index.IndexReader.CacheHelper;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.BytesRef;
+
+/**
+ * Class to do efficient primary-key lookups. Instances are not thread-safe. Each calling thread
+ * must own its own instance, and an instance must only ever be used from the thread that
+ * constructed it. When assertions are enabled this invariant is checked on every lookup and reopen
+ * call. Do not use this if a term may appear in more than one document! It will only return the
+ * first one it finds.
+ *
+ * @lucene.experimental
+ */
+public class PrimaryKeyLookup {
+
+  /** Name of the primary-key field. */
+  protected final String idFieldName;
+
+  /** Per-segment {@link TermsEnum}. */
+  protected final TermsEnum[] termsEnums;
+
+  /** Per-segment {@link PostingsEnum}. */
+  protected final PostingsEnum[] postingsEnums;
+
+  /** Per-segment live-docs bitset. */
+  protected final Bits[] liveDocs;
+
+  /** Per-segment {@code docBase} used to translate leaf docIDs to absolute docIDs. */
+  protected final int[] docBases;
+
+  /** Number of populated entries in termsEnums/postingsEnums. */
+  protected final int numEnums;
+
+  /** True if any segment in this lookup has deletions. */
+  protected final boolean hasDeletions;
+
+  /**
+   * Maps a segment core cache key to its index in termsEnums/postingsEnums, so that reopen can
+   * transfer per-segment enum state to the next generation.
+   */
+  protected final Map<IndexReader.CacheKey, Integer> enumIndexes;
+
+  // Thread-stickiness guard.
+  private final long ownerThreadId;
+
+  /** Construct a {@code PrimaryKeyLookup} bound to {@code reader}. */
+  public PrimaryKeyLookup(IndexReader reader, String idFieldName) throws IOException {
+    this(reader, idFieldName, Collections.emptyMap(), null, null);
+  }
+
+  /** Construct a {@code PrimaryKeyLookup} instance. */
+  protected PrimaryKeyLookup(
+      IndexReader reader,
+      String idFieldName,
+      Map<IndexReader.CacheKey, Integer> prevEnumIndexes,
+      TermsEnum[] reusableTermsEnums,
+      PostingsEnum[] reusablePostingsEnums)
+      throws IOException {
+    this.idFieldName = idFieldName;
+    this.ownerThreadId = Thread.currentThread().threadId();
+
+    List<LeafReaderContext> leaves = new ArrayList<>(reader.leaves());
+    // Larger segments are more likely to have the id, so we sort largest to smallest by numDocs:
+    leaves.sort((c1, c2) -> c2.reader().numDocs() - c1.reader().numDocs());
+
+    termsEnums = new TermsEnum[leaves.size()];
+    postingsEnums = new PostingsEnum[leaves.size()];
+    liveDocs = new Bits[leaves.size()];
+    docBases = new int[leaves.size()];
+    enumIndexes = new HashMap<>();
+    int numEnums = 0;
+    boolean hasDeletions = false;
+
+    for (LeafReaderContext context : leaves) {
+      LeafReader leafReader = context.reader();
+      CacheHelper cacheHelper = leafReader.getCoreCacheHelper();
+      IndexReader.CacheKey cacheKey = cacheHelper == null ? null : cacheHelper.getKey();
+
+      if (cacheKey != null && prevEnumIndexes.containsKey(cacheKey)) {
+        // Reuse termsEnum, postingsEnum.
+        int seg = prevEnumIndexes.get(cacheKey);
+        termsEnums[numEnums] = reusableTermsEnums[seg];
+        postingsEnums[numEnums] = reusablePostingsEnums[seg];
+      } else {
+        // New or empty segment.
+        Terms terms = leafReader.terms(idFieldName);
+        if (terms != null) {
+          termsEnums[numEnums] = terms.iterator();
+          assert termsEnums[numEnums] != null;
+        }
+      }
+
+      if (termsEnums[numEnums] != null) {
+        if (cacheKey != null) {
+          enumIndexes.put(cacheKey, numEnums);
+        }
+
+        docBases[numEnums] = context.docBase;
+        liveDocs[numEnums] = leafReader.getLiveDocs();
+        hasDeletions |= leafReader.hasDeletions();
+
+        numEnums++;
+      }
+    }
+
+    this.numEnums = numEnums;
+    this.hasDeletions = hasDeletions;
+  }
+
+  /** Returns docID if found, else -1. */
+  public int lookup(BytesRef id) throws IOException {
+    assert assertOwnerThread();
+    for (int seg = 0; seg < numEnums; seg++) {
+      if (termsEnums[seg].seekExact(id) == false) {
+        continue;
+      }
+      postingsEnums[seg] = termsEnums[seg].postings(postingsEnums[seg], 0);
+
+      if (liveDocs[seg] == null) {
+        // Fast path: segment has no deletions.
+        int docID = postingsEnums[seg].nextDoc();
+        assert docID != PostingsEnum.NO_MORE_DOCS;
+        return docBases[seg] + docID;
+      }
+
+      // Slow path: segment has deletions; scan until we find a live doc or exhaust.
+      int docID;
+      while ((docID = postingsEnums[seg].nextDoc()) != PostingsEnum.NO_MORE_DOCS) {
+        if (liveDocs[seg].get(docID)) {
+          return docBases[seg] + docID;
+        }
+      }
+
+      assert hasDeletions;
+    }
+    return -1;
+  }
+
+  /**
+   * Returns a new PrimaryKeyLookup instance bound reusing this instance per-segment
+   * TermsEnum/PostingsEnum etc.
+   */
+  public PrimaryKeyLookup reopen(IndexReader reader) throws IOException {
+    assert assertOwnerThread();
+    if (reader == null) {
+      return null;
+    }
+    return new PrimaryKeyLookup(
+        reader, this.idFieldName, this.enumIndexes, this.termsEnums, this.postingsEnums);
+  }
+
+  private boolean assertOwnerThread() {
+    long current = Thread.currentThread().threadId();
+    assert current == ownerThreadId
+        : "PrimaryKeyLookup is not thread-safe: it was created by thread "
+            + ownerThreadId
+            + " but accessed from thread "
+            + current;
+    return true;
+  }
+}

lucene/core/src/test/org/apache/lucene/index/TestPrimaryKeyLookup.java

@@ -0,0 +1,196 @@
+/*
+ * 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.lucene.index;
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.KeywordField;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.analysis.MockAnalyzer;
+import org.apache.lucene.tests.util.LuceneTestCase;
+
+public class TestPrimaryKeyLookup extends LuceneTestCase {
+
+  public void testReopen() throws Exception {
+    Directory dir = newDirectory();
+    IndexWriter writer =
+        new IndexWriter(
+            dir,
+            new IndexWriterConfig(new MockAnalyzer(random()))
+                .setMergePolicy(NoMergePolicy.INSTANCE));
+
+    Document doc;
+    doc = new Document();
+    doc.add(new KeywordField("PK", "1", Field.Store.NO));
+    writer.addDocument(doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "2", Field.Store.NO));
+    writer.addDocument(doc);
+    writer.flush();
+
+    // Segment with no PK terms.
+    doc = new Document();
+    doc.add(new KeywordField("PK2", "3", Field.Store.NO));
+    writer.addDocument(doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK2", "4", Field.Store.NO));
+    writer.addDocument(doc);
+    writer.flush();
+
+    DirectoryReader reader1 = DirectoryReader.open(writer);
+    PrimaryKeyLookup pk1 = new PrimaryKeyLookup(reader1, "PK");
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "5", Field.Store.NO));
+    writer.addDocument(doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "6", Field.Store.NO));
+    writer.addDocument(doc);
+    writer.deleteDocuments(new Term("PK", "1"));
+    writer.flush();
+
+    doc = new Document();
+    doc.add(new KeywordField("PK2", "7", Field.Store.NO));
+    writer.addDocument(doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK2", "8", Field.Store.NO));
+    writer.addDocument(doc);
+    writer.flush();
+
+    assertEquals(0, pk1.lookup(newBytesRef("1")));
+    assertEquals(1, pk1.lookup(newBytesRef("2")));
+    assertEquals(-1, pk1.lookup(newBytesRef("5")));
+    assertEquals(-1, pk1.lookup(newBytesRef("8")));
+
+    DirectoryReader reader2 = DirectoryReader.openIfChanged(reader1);
+    PrimaryKeyLookup pk2 = pk1.reopen(reader2);
+
+    assertEquals(-1, pk2.lookup(newBytesRef("1")));
+    assertEquals(1, pk2.lookup(newBytesRef("2")));
+    assertEquals(4, pk2.lookup(newBytesRef("5")));
+    assertEquals(-1, pk2.lookup(newBytesRef("8")));
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "9", Field.Store.NO));
+    writer.addDocument(doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "10", Field.Store.NO));
+    writer.addDocument(doc);
+    writer.flush();
+
+    assertEquals(-1, pk2.lookup(newBytesRef("9")));
+    DirectoryReader reader3 = DirectoryReader.openIfChanged(reader2);
+    PrimaryKeyLookup pk3 = pk2.reopen(reader3);
+    assertEquals(8, pk3.lookup(newBytesRef("9")));
+
+    DirectoryReader reader4 = DirectoryReader.openIfChanged(reader3);
+    assertNull(pk3.reopen(reader4));
+
+    writer.close();
+    reader1.close();
+    reader2.close();
+    reader3.close();
+    dir.close();
+  }
+
+  public void testPKLookupWithUpdate() throws Exception {
+    Directory dir = newDirectory();
+    IndexWriter writer =
+        new IndexWriter(
+            dir,
+            new IndexWriterConfig(new MockAnalyzer(random()))
+                .setMergePolicy(NoMergePolicy.INSTANCE));
+
+    Document doc;
+    doc = new Document();
+    doc.add(new KeywordField("PK", "1", Field.Store.NO));
+    doc.add(new KeywordField("version", "1", Field.Store.NO));
+    writer.addDocument(doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "1", Field.Store.NO));
+    doc.add(new KeywordField("version", "2", Field.Store.NO));
+    writer.updateDocument(new Term("PK", "1"), doc);
+
+    doc = new Document();
+    doc.add(new KeywordField("PK", "1", Field.Store.NO));
+    doc.add(new KeywordField("version", "3", Field.Store.NO));
+    writer.updateDocument(new Term("PK", "1"), doc);
+    writer.flush();
+    writer.close();
+
+    DirectoryReader reader = DirectoryReader.open(dir);
+    PrimaryKeyLookup pk = new PrimaryKeyLookup(reader, "PK");
+    assertEquals(2, pk.lookup(newBytesRef("1")));
+
+    reader.close();
+    dir.close();
+  }
+
+  /**
+   * Constructing on one thread and calling {@link PrimaryKeyLookup#lookup} from another must fire
+   * the thread-stickiness assertion.
+   */
+  public void testThreadStickinessAssertion() throws Exception {
+    Directory dir = newDirectory();
+    IndexWriter writer = new IndexWriter(dir, new IndexWriterConfig(new MockAnalyzer(random())));
+    Document doc = new Document();
+    doc.add(new KeywordField("PK", "1", Field.Store.NO));
+    writer.addDocument(doc);
+    writer.commit();
+    writer.close();
+
+    DirectoryReader reader = DirectoryReader.open(dir);
+    PrimaryKeyLookup pk = new PrimaryKeyLookup(reader, "PK");
+
+    if (TEST_ASSERTS_ENABLED) {
+      // pk was constructed on this thread; calling lookup from any other thread must trip
+      // the thread-stickiness assertion in PrimaryKeyLookup.
+      class CrossThreadLookup extends Thread {
+        Throwable thrown;
+
+        @Override
+        public void run() {
+          try {
+            pk.lookup(newBytesRef("1"));
+          } catch (Throwable t) {
+            thrown = t;
+          }
+        }
+      }
+
+      CrossThreadLookup other = new CrossThreadLookup();
+      other.setName("pk-cross-thread-lookup");
+      other.start();
+      other.join();
+
+      assertTrue(
+          "expected AssertionError but got: " + other.thrown,
+          other.thrown instanceof AssertionError);
+      assertTrue(
+          "unexpected message: " + other.thrown.getMessage(),
+          other.thrown.getMessage().contains("not thread-safe"));
+    }
+    reader.close();
+    dir.close();
+  }
+}