Add a hook to filter out documents by query to IndexModule

This is a more limited version of the existing ReaderWrapper
that only accepts a query of documents to filter from the
reader. It is meant to support temporarily filtering out
documents that no longer belong to a shard after it has
been split, until the shard is cleaned.

The decision about whether to install the filter is made
when the index module is created. To reduce overhead in
the common case, the query-supplying function may return
null in which case the underlying reader is returned
directly.

Author: bcully

server/src/main/java/org/elasticsearch/common/lucene/Lucene.java

@@ -29,12 +29,14 @@
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.IndexFormatTooNewException;
 import org.apache.lucene.index.IndexFormatTooOldException;
+import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.NoMergePolicy;
 import org.apache.lucene.index.NoMergeScheduler;
+import org.apache.lucene.index.ReaderUtil;
 import org.apache.lucene.index.SegmentCommitInfo;
 import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.index.SegmentReader;
@@ -60,9 +62,11 @@
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.Lock;
+import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.Version;
+import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
@@ -78,6 +82,8 @@
 import org.elasticsearch.index.analysis.NamedAnalyzer;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.lucene.grouping.TopFieldGroups;
+import org.elasticsearch.lucene.util.BitSets;
+import org.elasticsearch.lucene.util.MatchAllBitSet;
 import org.elasticsearch.search.sort.ShardDocSortField;
 
 import java.io.IOException;
@@ -1013,6 +1019,168 @@ public CacheHelper getReaderCacheHelper() {
         }
     }
 
+    public static DirectoryReader queryFilteredDirectoryReader(DirectoryReader in, Query query) throws IOException {
+        return new QueryFilterDirectoryReader(in, query);
+    }
+
+    private static class QueryFilterDirectoryReader extends FilterDirectoryReader {
+        private final Query query;
+
+        QueryFilterDirectoryReader(DirectoryReader in, Query query) throws IOException {
+            super(in, new SubReaderWrapper() {
+                @Override
+                public LeafReader wrap(LeafReader reader) {
+                    return new QueryFilterLeafReader(reader, query);
+                }
+            });
+            this.query = query;
+        }
+
+        @Override
+        protected DirectoryReader doWrapDirectoryReader(DirectoryReader in) throws IOException {
+            return new QueryFilterDirectoryReader(in, query);
+        }
+
+        @Override
+        public CacheHelper getReaderCacheHelper() {
+            return in.getReaderCacheHelper();
+        }
+    }
+
+    private static class QueryFilterLeafReader extends SequentialStoredFieldsLeafReader {
+        private final Query query;
+
+        private int numDocs = -1;
+        private BitSet unownedDocs;
+
+        protected QueryFilterLeafReader(LeafReader in, Query query) {
+            super(in);
+            this.query = query;
+        }
+
+        /**
+         * Returns all documents that are not deleted and are owned by the current shard.
+         * We need to recalculate this every time because `in.getLiveDocs()` can change when deletes are performed.
+         */
+        @Override
+        public Bits getLiveDocs() {
+            ensureUnownedDocumentsPresent();
+            Bits actualLiveDocs = in.getLiveDocs();
+
+            if (unownedDocs == null) {
+                return actualLiveDocs;
+            }
+
+            if (unownedDocs instanceof MatchAllBitSet) {
+                return new Bits.MatchNoBits(in.maxDoc());
+            }
+
+            var liveDocsBitsWithAllLiveCheck = actualLiveDocs == null ? new MatchAllBitSet(in.maxDoc()) : actualLiveDocs;
+            return new FilterBits(liveDocsBitsWithAllLiveCheck, unownedDocs);
+        }
+
+        @Override
+        public int numDocs() {
+            ensureUnownedDocumentsPresent();
+            return numDocs;
+        }
+
+        @Override
+        public boolean hasDeletions() {
+            // It is possible that there are unowned docs which we are going to present as deletes.
+            return true;
+        }
+
+        @Override
+        protected StoredFieldsReader doGetSequentialStoredFieldsReader(StoredFieldsReader reader) {
+            return reader;
+        }
+
+        @Override
+        public CacheHelper getCoreCacheHelper() {
+            return in.getCoreCacheHelper();
+        }
+
+        @Override
+        public CacheHelper getReaderCacheHelper() {
+            // Not delegated since we change live docs.
+            return null;
+        }
+
+        private void ensureUnownedDocumentsPresent() {
+            if (numDocs == -1) {
+                synchronized (this) {
+                    if (numDocs == -1) {
+                        try {
+                            unownedDocs = queryUnownedDocs();
+                            numDocs = calculateNumDocs(in, unownedDocs);
+                        } catch (Exception e) {
+                            throw new ElasticsearchException("Failed to execute unowned documents query", e);
+                        }
+                    }
+                }
+            }
+        }
+
+        private BitSet queryUnownedDocs() throws IOException {
+            final IndexReaderContext topLevelContext = ReaderUtil.getTopLevelContext(in.getContext());
+
+            final IndexSearcher searcher = new IndexSearcher(topLevelContext);
+            searcher.setQueryCache(null);
+
+            final Query rewrittenQuery = searcher.rewrite(query);
+            // TODO there is a possible optimization of checking for MatchAllDocsQuery which would mean that all documents are unowned.
+            final Weight weight = searcher.createWeight(rewrittenQuery, ScoreMode.COMPLETE_NO_SCORES, 1f);
+            final Scorer s = weight.scorer(in.getContext());
+            if (s == null) {
+                return null;
+            } else {
+                return BitSets.of(s.iterator(), in.maxDoc());
+            }
+        }
+
+        private static int calculateNumDocs(LeafReader reader, BitSet unownedDocs) {
+            final Bits liveDocs = reader.getLiveDocs();
+
+            // No deleted documents are present, therefore number of documents is total minus unowned.
+            if (liveDocs == null) {
+                return reader.numDocs() - unownedDocs.cardinality();
+            }
+
+            if (unownedDocs instanceof MatchAllBitSet) {
+                return 0;
+            }
+
+            int numDocs = 0;
+            for (int i = 0; i < liveDocs.length(); i++) {
+                if (liveDocs.get(i) && unownedDocs.get(i) == false) {
+                    numDocs++;
+                }
+            }
+            return numDocs;
+        }
+    }
+
+    static class FilterBits implements Bits {
+        private final Bits original;
+        private final Bits filteredOut;
+
+        FilterBits(Bits original, BitSet filteredOut) {
+            this.original = original;
+            this.filteredOut = filteredOut;
+        }
+
+        @Override
+        public boolean get(int index) {
+            return original.get(index) && (filteredOut.get(index) == false);
+        }
+
+        @Override
+        public int length() {
+            return original.length();
+        }
+    }
+
     private static int popCount(Bits bits) {
         assert bits != null;
         int onBits = 0;

server/src/main/java/org/elasticsearch/index/IndexModule.java

@@ -15,6 +15,7 @@
 import org.apache.lucene.index.FilterDirectoryReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReader;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.similarities.BM25Similarity;
 import org.apache.lucene.search.similarities.Similarity;
 import org.apache.lucene.store.Directory;
@@ -26,6 +27,7 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.TriFunction;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Setting.Property;
 import org.elasticsearch.common.settings.Settings;
@@ -168,6 +170,9 @@ public interface DirectoryWrapper {
     private final SetOnce<DirectoryWrapper> indexDirectoryWrapper = new SetOnce<>();
     private final SetOnce<Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>>> indexReaderWrapper =
         new SetOnce<>();
+    // a limited version of the reader wrapper that only returns a query that can filter out documents
+    private final SetOnce<Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>>> indexReaderFilter =
+        new SetOnce<>();
     private final Set<IndexEventListener> indexEventListeners = new HashSet<>();
     private final Map<String, TriFunction<Settings, IndexVersion, ScriptService, Similarity>> similarities = new HashMap<>();
     private final Map<String, IndexStorePlugin.DirectoryFactory> directoryFactories;
@@ -391,6 +396,29 @@ public void setReaderWrapper(
         this.indexReaderWrapper.set(indexReaderWrapperFactory);
     }
 
+    /**
+     * A limited version of {@link IndexModule#setReaderWrapper} that can only filter out documents matching a query.
+     * <p>
+     * The returned query may be null in which case no documents are filtered.
+     * @param indexReaderFilterFactory the factory for creating per-DirectoryReader @link Query} instances
+     */
+    public void setReaderFilter(Function<IndexService, CheckedFunction<DirectoryReader, Query, IOException>> indexReaderFilterFactory) {
+        ensureNotFrozen();
+        // convert the query filter into a reader wrapper
+        Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = (indexService) -> {
+            var queryFactory = indexReaderFilterFactory.apply(indexService);
+            return (reader) -> {
+                Query query = queryFactory.apply(reader);
+                if (query == null) {
+                    return reader;
+                }
+                return Lucene.queryFilteredDirectoryReader(reader, query);
+            };
+        };
+
+        this.indexReaderFilter.set(readerWrapperFactory);
+    }
+
     /**
      * Sets a {@link Directory} wrapping method that allows to apply a function to the Lucene directory instance
      * created by {@link org.elasticsearch.plugins.IndexStorePlugin.DirectoryFactory}.
@@ -499,8 +527,23 @@ public IndexService newIndexService(
         QueryRewriteInterceptor queryRewriteInterceptor
     ) throws IOException {
         final IndexEventListener eventListener = freeze();
-        Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrapperFactory = indexReaderWrapper
-            .get() == null ? (shard) -> null : indexReaderWrapper.get();
+        var readerFilterFactory = indexReaderFilter.get();
+        var readerWrapperFactory = indexReaderWrapper.get();
+
+        Function<IndexService, CheckedFunction<DirectoryReader, DirectoryReader, IOException>> readerWrappersFactory = (indexService) -> {
+            if (readerFilterFactory == null && readerWrapperFactory == null) {
+                return null;
+            } else if (readerFilterFactory == null) {
+                return readerWrapperFactory.apply(indexService);
+            } else if (readerWrapperFactory == null) {
+                return readerFilterFactory.apply(indexService);
+            } else {
+                var readerFilter = readerFilterFactory.apply(indexService);
+                var readerWrapper = readerWrapperFactory.apply(indexService);
+                return (reader) -> readerWrapper.apply(readerFilter.apply(reader));
+            }
+        };
+
         eventListener.beforeIndexCreated(indexSettings.getIndex(), indexSettings.getSettings());
         final IndexStorePlugin.DirectoryFactory directoryFactory = getDirectoryFactory(indexSettings, directoryFactories);
         final IndexStorePlugin.RecoveryStateFactory recoveryStateFactory = getRecoveryStateFactory(indexSettings, recoveryStateFactories);
@@ -545,7 +588,7 @@ public IndexService newIndexService(
                 queryCache,
                 directoryFactory,
                 eventListener,
-                readerWrapperFactory,
+                readerWrappersFactory,
                 mapperRegistry,
                 indicesFieldDataCache,
                 searchOperationListeners,

server/src/test/java/org/elasticsearch/common/lucene/LuceneTests.java

@@ -550,6 +550,32 @@ public void testWrapLiveDocsNotExposeAbortedDocuments() throws Exception {
         IOUtils.close(writer, dir);
     }
 
+    public void testQueryFilteredDirectoryReader() throws IOException {
+        try (Directory dir = newDirectory()) {
+            try (RandomIndexWriter w = new RandomIndexWriter(random(), dir)) {
+
+                Document doc = new Document();
+                doc.add(new StringField("foo", "bar", Store.YES));
+                w.addDocument(doc);
+
+                doc = new Document();
+                doc.add(new StringField("foo", "baz", Store.YES));
+                w.addDocument(doc);
+
+                try (DirectoryReader reader = w.getReader()) {
+                    assertThat(reader.numDocs(), equalTo(2));
+                    DirectoryReader filtered = Lucene.queryFilteredDirectoryReader(reader, new TermQuery(new Term("foo", "bar")));
+                    assertThat(filtered.numDocs(), equalTo(1));
+                    IndexSearcher searcher = newSearcher(filtered);
+                    TopDocs topDocs = searcher.search(new MatchAllDocsQuery(), Integer.MAX_VALUE);
+                    StoredFields storedFields = filtered.storedFields();
+                    assertThat(topDocs.totalHits.value(), equalTo(1L));
+                    assertThat(storedFields.document(topDocs.scoreDocs[0].doc).get("foo"), equalTo("baz"));
+                }
+            }
+        }
+    }
+
     public void testSortFieldSerialization() throws IOException {
         Tuple<SortField, SortField> sortFieldTuple = randomSortField();
         SortField deserialized = copyInstance(