org.elasticsearch.index.shard.ShardSplittingQuery Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch Show documentation
Show all versions of elasticsearch Show documentation
Elasticsearch subproject :server
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/
package org.elasticsearch.index.shard;
import org.apache.lucene.index.FieldInfo;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.PostingsEnum;
import org.apache.lucene.index.StoredFieldVisitor;
import org.apache.lucene.index.StoredFields;
import org.apache.lucene.index.Terms;
import org.apache.lucene.index.TermsEnum;
import org.apache.lucene.search.ConstantScoreScorer;
import org.apache.lucene.search.ConstantScoreWeight;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.QueryVisitor;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.TwoPhaseIterator;
import org.apache.lucene.search.Weight;
import org.apache.lucene.search.join.BitSetProducer;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.BitSetIterator;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.FixedBitSet;
import org.elasticsearch.cluster.metadata.IndexMetadata;
import org.elasticsearch.cluster.routing.IndexRouting;
import org.elasticsearch.common.lucene.search.Queries;
import org.elasticsearch.index.IndexVersion;
import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
import org.elasticsearch.index.mapper.IdFieldMapper;
import org.elasticsearch.index.mapper.RoutingFieldMapper;
import org.elasticsearch.index.mapper.Uid;
import java.io.IOException;
import java.util.function.Function;
import java.util.function.IntConsumer;
import java.util.function.Predicate;
/**
* A query that selects all docs that do NOT belong in the current shards this query is executed on.
* It can be used to split a shard into N shards marking every document that doesn't belong into the shard
* as deleted. See {@link org.apache.lucene.index.IndexWriter#deleteDocuments(Query...)}
*/
final class ShardSplittingQuery extends Query {
private final IndexMetadata indexMetadata;
private final IndexRouting indexRouting;
private final int shardId;
private final BitSetProducer nestedParentBitSetProducer;
ShardSplittingQuery(IndexMetadata indexMetadata, int shardId, boolean hasNested) {
this.indexMetadata = indexMetadata;
this.indexRouting = IndexRouting.fromIndexMetadata(indexMetadata);
this.shardId = shardId;
this.nestedParentBitSetProducer = hasNested ? newParentDocBitSetProducer(indexMetadata.getCreationVersion()) : null;
}
@Override
public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) {
return new ConstantScoreWeight(this, boost) {
@Override
public String toString() {
return "weight(delete docs query)";
}
@Override
public Scorer scorer(LeafReaderContext context) throws IOException {
LeafReader leafReader = context.reader();
FixedBitSet bitSet = new FixedBitSet(leafReader.maxDoc());
Terms terms = leafReader.terms(RoutingFieldMapper.NAME);
Predicate includeInShard = ref -> {
// TODO IndexRouting should build the query somehow
int targetShardId = indexRouting.getShard(Uid.decodeId(ref.bytes, ref.offset, ref.length), null);
return shardId == targetShardId;
};
if (terms == null) {
// this is the common case - no partitioning and no _routing values
// in this case we also don't do anything special with regards to nested docs since we basically delete
// by ID and parent and nested all have the same id.
assert indexMetadata.isRoutingPartitionedIndex() == false;
findSplitDocs(IdFieldMapper.NAME, includeInShard, leafReader, bitSet::set);
} else {
final BitSet parentBitSet;
if (nestedParentBitSetProducer == null) {
parentBitSet = null;
} else {
parentBitSet = nestedParentBitSetProducer.getBitSet(context);
if (parentBitSet == null) {
return null; // no matches
}
}
if (indexMetadata.isRoutingPartitionedIndex()) {
// this is the heaviest invariant. Here we have to visit all docs stored fields do extract _id and _routing
// this index is routing partitioned.
Visitor visitor = new Visitor(leafReader);
TwoPhaseIterator twoPhaseIterator = parentBitSet == null
? new RoutingPartitionedDocIdSetIterator(visitor)
: new NestedRoutingPartitionedDocIdSetIterator(visitor, parentBitSet);
return new ConstantScoreScorer(this, score(), scoreMode, twoPhaseIterator);
} else {
// here we potentially guard the docID consumers with our parent bitset if we have one.
// this ensures that we are only marking root documents in the nested case and if necessary
// we do a second pass to mark the corresponding children in markChildDocs
Function maybeWrapConsumer = consumer -> {
if (parentBitSet != null) {
return docId -> {
if (parentBitSet.get(docId)) {
consumer.accept(docId);
}
};
}
return consumer;
};
// in the _routing case we first go and find all docs that have a routing value and mark the ones we have to delete
findSplitDocs(RoutingFieldMapper.NAME, ref -> {
int targetShardId = indexRouting.getShard(null, ref.utf8ToString());
return shardId == targetShardId;
}, leafReader, maybeWrapConsumer.apply(bitSet::set));
// TODO have the IndexRouting build the query and pass routingRequired in
boolean routingRequired = indexMetadata.mapping() == null ? false : indexMetadata.mapping().routingRequired();
// now if we have a mixed index where some docs have a _routing value and some don't we have to exclude the ones
// with a routing value from the next iteration and delete / select based on the ID.
if (routingRequired == false && terms.getDocCount() != leafReader.maxDoc()) {
/*
* This is a special case where some docs don't have routing values.
* It's annoying, but it's allowed to build an index where some documents
* hve routing and others don't.
*
* Luckily, if the routing field is required in the mapping then we can
* safely assume that all documents which are don't have a routing are
* nested documents. And we pick those up later based on the assignment
* of the document that contains them.
*/
FixedBitSet hasRoutingValue = new FixedBitSet(leafReader.maxDoc());
findSplitDocs(RoutingFieldMapper.NAME, ref -> false, leafReader, maybeWrapConsumer.apply(hasRoutingValue::set));
IntConsumer bitSetConsumer = maybeWrapConsumer.apply(bitSet::set);
findSplitDocs(IdFieldMapper.NAME, includeInShard, leafReader, docId -> {
if (hasRoutingValue.get(docId) == false) {
bitSetConsumer.accept(docId);
}
});
}
}
if (parentBitSet != null) {
// if nested docs are involved we also need to mark all child docs that belong to a matching parent doc.
markChildDocs(parentBitSet, bitSet);
}
}
return new ConstantScoreScorer(this, score(), scoreMode, new BitSetIterator(bitSet, bitSet.length()));
}
@Override
public boolean isCacheable(LeafReaderContext ctx) {
// This is not a regular query, let's not cache it. It wouldn't help
// anyway.
return false;
}
};
}
private static void markChildDocs(BitSet parentDocs, BitSet matchingDocs) {
int currentDeleted = 0;
while (currentDeleted < matchingDocs.length()
&& (currentDeleted = matchingDocs.nextSetBit(currentDeleted)) != DocIdSetIterator.NO_MORE_DOCS) {
int previousParent = parentDocs.prevSetBit(Math.max(0, currentDeleted - 1));
for (int i = previousParent + 1; i < currentDeleted; i++) {
matchingDocs.set(i);
}
currentDeleted++;
}
}
@Override
public String toString(String field) {
return "shard_splitting_query";
}
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (sameClassAs(o) == false) return false;
ShardSplittingQuery that = (ShardSplittingQuery) o;
if (shardId != that.shardId) return false;
return indexMetadata.equals(that.indexMetadata);
}
@Override
public int hashCode() {
int result = indexMetadata.hashCode();
result = 31 * result + shardId;
return classHash() ^ result;
}
@Override
public void visit(QueryVisitor visitor) {
visitor.visitLeaf(this);
}
private static void findSplitDocs(String idField, Predicate includeInShard, LeafReader leafReader, IntConsumer consumer)
throws IOException {
Terms terms = leafReader.terms(idField);
TermsEnum iterator = terms.iterator();
BytesRef idTerm;
PostingsEnum postingsEnum = null;
while ((idTerm = iterator.next()) != null) {
if (includeInShard.test(idTerm) == false) {
postingsEnum = iterator.postings(postingsEnum);
int doc;
while ((doc = postingsEnum.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
consumer.accept(doc);
}
}
}
}
/* this class is a stored fields visitor that reads _id and/or _routing from the stored fields which is necessary in the case
of a routing partitioned index sine otherwise we would need to un-invert the _id and _routing field which is memory heavy */
private final class Visitor extends StoredFieldVisitor {
final LeafReader leafReader;
final StoredFields storedFields;
private int leftToVisit = 2;
private String routing;
private String id;
Visitor(LeafReader leafReader) throws IOException {
this.leafReader = leafReader;
this.storedFields = leafReader.storedFields();
}
@Override
public void binaryField(FieldInfo fieldInfo, byte[] value) throws IOException {
switch (fieldInfo.name) {
case IdFieldMapper.NAME -> id = Uid.decodeId(value);
default -> throw new IllegalStateException("Unexpected field: " + fieldInfo.name);
}
}
@Override
public void stringField(FieldInfo fieldInfo, String value) throws IOException {
switch (fieldInfo.name) {
case RoutingFieldMapper.NAME -> routing = value;
default -> throw new IllegalStateException("Unexpected field: " + fieldInfo.name);
}
}
@Override
public Status needsField(FieldInfo fieldInfo) throws IOException {
// we don't support 5.x so no need for the uid field
switch (fieldInfo.name) {
case IdFieldMapper.NAME:
case RoutingFieldMapper.NAME:
leftToVisit--;
return Status.YES;
default:
return leftToVisit == 0 ? Status.STOP : Status.NO;
}
}
boolean matches(int doc) throws IOException {
routing = id = null;
leftToVisit = 2;
storedFields.document(doc, this);
assert id != null : "docID must not be null - we might have hit a nested document";
int targetShardId = indexRouting.getShard(id, routing);
return targetShardId != shardId;
}
}
/**
* This two phase iterator visits every live doc and selects all docs that don't belong into this
* shard based on their id and routing value. This is only used in a routing partitioned index.
*/
private static final class RoutingPartitionedDocIdSetIterator extends TwoPhaseIterator {
private final Visitor visitor;
RoutingPartitionedDocIdSetIterator(Visitor visitor) {
super(DocIdSetIterator.all(visitor.leafReader.maxDoc())); // we iterate all live-docs
this.visitor = visitor;
}
@Override
public boolean matches() throws IOException {
return visitor.matches(approximation.docID());
}
@Override
public float matchCost() {
return 42; // that's obvious, right?
}
}
/**
* This TwoPhaseIterator marks all nested docs of matching parents as matches as well.
*/
private static final class NestedRoutingPartitionedDocIdSetIterator extends TwoPhaseIterator {
private final Visitor visitor;
private final BitSet parentDocs;
private int nextParent = -1;
private boolean nextParentMatches;
NestedRoutingPartitionedDocIdSetIterator(Visitor visitor, BitSet parentDocs) {
super(DocIdSetIterator.all(visitor.leafReader.maxDoc())); // we iterate all live-docs
this.parentDocs = parentDocs;
this.visitor = visitor;
}
@Override
public boolean matches() throws IOException {
// the educated reader might ask why this works, it does because all live doc ids (root docs and nested docs) are evaluated in
// order and that way we don't need to seek backwards as we do in other nested docs cases.
int doc = approximation.docID();
if (doc > nextParent) {
// we only check once per nested/parent set
nextParent = parentDocs.nextSetBit(doc);
// never check a child document against the visitor, they neihter have _id nor _routing as stored fields
nextParentMatches = visitor.matches(nextParent);
}
return nextParentMatches;
}
@Override
public float matchCost() {
return 42; // that's obvious, right?
}
}
/*
* this is used internally to obtain a bitset for parent documents. We don't cache this since we never access the same reader more
* than once. There is no point in using BitsetFilterCache#BitSetProducerWarmer since we use this only as a delete by query which is
* executed on a recovery-private index writer. There is no point in caching it and it won't have a cache hit either.
*/
private static BitSetProducer newParentDocBitSetProducer(IndexVersion indexCreationVersion) {
return context -> BitsetFilterCache.bitsetFromQuery(Queries.newNonNestedFilter(indexCreationVersion), context);
}
}