From 0a4245552422db90a8ea431983e33c13dd9463f3 Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Mon, 12 Aug 2024 09:02:33 +0900 Subject: [PATCH] Speed up dense/sparse vector stats (#111729) This change ensures that we don't try to compute stats on mappings that don't have dense or sparse vector fields. We don't need to go through all the fields on every segment, instead we can extract the vector fields upfront and limit the work to only indices that define these types. Closes #111715 --- docs/changelog/111729.yaml | 6 ++ .../elasticsearch/index/engine/Engine.java | 82 +++++++++++-------- .../elasticsearch/index/shard/IndexShard.java | 3 +- .../index/engine/frozen/FrozenEngine.java | 14 ++-- 4 files changed, 61 insertions(+), 44 deletions(-) create mode 100644 docs/changelog/111729.yaml diff --git a/docs/changelog/111729.yaml b/docs/changelog/111729.yaml new file mode 100644 index 0000000000000..c75c14a997da9 --- /dev/null +++ b/docs/changelog/111729.yaml @@ -0,0 +1,6 @@ +pr: 111729 +summary: Speed up dense/sparse vector stats +area: Vector Search +type: bug +issues: + - 111715 diff --git a/server/src/main/java/org/elasticsearch/index/engine/Engine.java b/server/src/main/java/org/elasticsearch/index/engine/Engine.java index 6f4511483126f..b07132eea75e8 100644 --- a/server/src/main/java/org/elasticsearch/index/engine/Engine.java +++ b/server/src/main/java/org/elasticsearch/index/engine/Engine.java @@ -61,7 +61,6 @@ import org.elasticsearch.index.IndexVersion; import org.elasticsearch.index.VersionType; import org.elasticsearch.index.mapper.DocumentParser; -import org.elasticsearch.index.mapper.FieldMapper; import org.elasticsearch.index.mapper.FieldNamesFieldMapper; import org.elasticsearch.index.mapper.LuceneDocument; import org.elasticsearch.index.mapper.Mapper; @@ -69,6 +68,7 @@ import org.elasticsearch.index.mapper.MappingLookup; import org.elasticsearch.index.mapper.ParsedDocument; import org.elasticsearch.index.mapper.Uid; +import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; import org.elasticsearch.index.mapper.vectors.SparseVectorFieldMapper; import org.elasticsearch.index.merge.MergeStats; import org.elasticsearch.index.seqno.SeqNoStats; @@ -242,19 +242,32 @@ protected final DocsStats docsStats(IndexReader indexReader) { /** * Returns the {@link DenseVectorStats} for this engine */ - public DenseVectorStats denseVectorStats() { + public DenseVectorStats denseVectorStats(MappingLookup mappingLookup) { + if (mappingLookup == null) { + return new DenseVectorStats(0); + } + + List fields = new ArrayList<>(); + for (Mapper mapper : mappingLookup.fieldMappers()) { + if (mapper instanceof DenseVectorFieldMapper) { + fields.add(mapper.fullPath()); + } + } + if (fields.isEmpty()) { + return new DenseVectorStats(0); + } try (Searcher searcher = acquireSearcher(DOC_STATS_SOURCE, SearcherScope.INTERNAL)) { - return denseVectorStats(searcher.getIndexReader()); + return denseVectorStats(searcher.getIndexReader(), fields); } } - protected final DenseVectorStats denseVectorStats(IndexReader indexReader) { + protected final DenseVectorStats denseVectorStats(IndexReader indexReader, List fields) { long valueCount = 0; // we don't wait for a pending refreshes here since it's a stats call instead we mark it as accessed only which will cause // the next scheduled refresh to go through and refresh the stats as well for (LeafReaderContext readerContext : indexReader.leaves()) { try { - valueCount += getDenseVectorValueCount(readerContext.reader()); + valueCount += getDenseVectorValueCount(readerContext.reader(), fields); } catch (IOException e) { logger.trace(() -> "failed to get dense vector stats for [" + readerContext + "]", e); } @@ -262,9 +275,10 @@ protected final DenseVectorStats denseVectorStats(IndexReader indexReader) { return new DenseVectorStats(valueCount); } - private long getDenseVectorValueCount(final LeafReader atomicReader) throws IOException { + private long getDenseVectorValueCount(final LeafReader atomicReader, List fields) throws IOException { long count = 0; - for (FieldInfo info : atomicReader.getFieldInfos()) { + for (var field : fields) { + var info = atomicReader.getFieldInfos().fieldInfo(field); if (info.getVectorDimension() > 0) { switch (info.getVectorEncoding()) { case FLOAT32 -> { @@ -285,23 +299,31 @@ private long getDenseVectorValueCount(final LeafReader atomicReader) throws IOEx * Returns the {@link SparseVectorStats} for this engine */ public SparseVectorStats sparseVectorStats(MappingLookup mappingLookup) { + if (mappingLookup == null) { + return new SparseVectorStats(0); + } + List fields = new ArrayList<>(); + for (Mapper mapper : mappingLookup.fieldMappers()) { + if (mapper instanceof SparseVectorFieldMapper) { + fields.add(new BytesRef(mapper.fullPath())); + } + } + if (fields.isEmpty()) { + return new SparseVectorStats(0); + } + Collections.sort(fields); try (Searcher searcher = acquireSearcher(DOC_STATS_SOURCE, SearcherScope.INTERNAL)) { - return sparseVectorStats(searcher.getIndexReader(), mappingLookup); + return sparseVectorStats(searcher.getIndexReader(), fields); } } - protected final SparseVectorStats sparseVectorStats(IndexReader indexReader, MappingLookup mappingLookup) { + protected final SparseVectorStats sparseVectorStats(IndexReader indexReader, List fields) { long valueCount = 0; - - if (mappingLookup == null) { - return new SparseVectorStats(valueCount); - } - // we don't wait for a pending refreshes here since it's a stats call instead we mark it as accessed only which will cause // the next scheduled refresh to go through and refresh the stats as well for (LeafReaderContext readerContext : indexReader.leaves()) { try { - valueCount += getSparseVectorValueCount(readerContext.reader(), mappingLookup); + valueCount += getSparseVectorValueCount(readerContext.reader(), fields); } catch (IOException e) { logger.trace(() -> "failed to get sparse vector stats for [" + readerContext + "]", e); } @@ -309,28 +331,16 @@ protected final SparseVectorStats sparseVectorStats(IndexReader indexReader, Map return new SparseVectorStats(valueCount); } - private long getSparseVectorValueCount(final LeafReader atomicReader, MappingLookup mappingLookup) throws IOException { + private long getSparseVectorValueCount(final LeafReader atomicReader, List fields) throws IOException { long count = 0; - - Map mappers = new HashMap<>(); - for (Mapper mapper : mappingLookup.fieldMappers()) { - if (mapper instanceof FieldMapper fieldMapper) { - if (fieldMapper.fieldType() instanceof SparseVectorFieldMapper.SparseVectorFieldType) { - mappers.put(fieldMapper.fullPath(), fieldMapper); - } - } - } - - for (FieldInfo info : atomicReader.getFieldInfos()) { - String name = info.name; - if (mappers.containsKey(name)) { - Terms terms = atomicReader.terms(FieldNamesFieldMapper.NAME); - if (terms != null) { - TermsEnum termsEnum = terms.iterator(); - if (termsEnum.seekExact(new BytesRef(name))) { - count += termsEnum.docFreq(); - } - } + Terms terms = atomicReader.terms(FieldNamesFieldMapper.NAME); + if (terms == null) { + return count; + } + TermsEnum termsEnum = terms.iterator(); + for (var fieldName : fields) { + if (termsEnum.seekExact(fieldName)) { + count += termsEnum.docFreq(); } } return count; diff --git a/server/src/main/java/org/elasticsearch/index/shard/IndexShard.java b/server/src/main/java/org/elasticsearch/index/shard/IndexShard.java index 881f4602be1c7..b15d3e08a0bb8 100644 --- a/server/src/main/java/org/elasticsearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/elasticsearch/index/shard/IndexShard.java @@ -1425,7 +1425,8 @@ public CompletionStats completionStats(String... fields) { public DenseVectorStats denseVectorStats() { readAllowed(); - return getEngine().denseVectorStats(); + MappingLookup mappingLookup = mapperService != null ? mapperService.mappingLookup() : null; + return getEngine().denseVectorStats(mappingLookup); } public SparseVectorStats sparseVectorStats() { diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/index/engine/frozen/FrozenEngine.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/index/engine/frozen/FrozenEngine.java index 0a13aab82aced..3b242ca94ac61 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/index/engine/frozen/FrozenEngine.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/index/engine/frozen/FrozenEngine.java @@ -63,8 +63,6 @@ public final class FrozenEngine extends ReadOnlyEngine { ); private final SegmentsStats segmentsStats; private final DocsStats docsStats; - private final DenseVectorStats denseVectorStats; - private final SparseVectorStats sparseVectorStats; private volatile ElasticsearchDirectoryReader lastOpenedReader; private final ElasticsearchDirectoryReader canMatchReader; private final Object cacheIdentity = new Object(); @@ -95,8 +93,6 @@ public FrozenEngine( fillSegmentStats(segmentReader, true, segmentsStats); } this.docsStats = docsStats(reader); - this.denseVectorStats = denseVectorStats(reader); - this.sparseVectorStats = sparseVectorStats(reader, null); canMatchReader = ElasticsearchDirectoryReader.wrap( new RewriteCachingDirectoryReader(directory, reader.leaves(), null), config.getShardId() @@ -334,13 +330,17 @@ public DocsStats docStats() { } @Override - public DenseVectorStats denseVectorStats() { - return denseVectorStats; + public DenseVectorStats denseVectorStats(MappingLookup mappingLookup) { + // We could cache the result on first call but dense vectors on frozen tier + // are very unlikely, so we just don't count them in the stats. + return new DenseVectorStats(0); } @Override public SparseVectorStats sparseVectorStats(MappingLookup mappingLookup) { - return sparseVectorStats; + // We could cache the result on first call but sparse vectors on frozen tier + // are very unlikely, so we just don't count them in the stats. + return new SparseVectorStats(0); } synchronized boolean isReaderOpen() {