From f0d50aebdc4c79b27d3857720db7c869896ef065 Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Fri, 9 Aug 2024 08:58:40 +0900 Subject: [PATCH] Speed up dense/sparse vector stats 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 --- .../elasticsearch/index/engine/Engine.java | 82 +++++++++++-------- .../elasticsearch/index/shard/IndexShard.java | 3 +- .../index/engine/frozen/FrozenEngine.java | 4 +- 3 files changed, 51 insertions(+), 38 deletions(-) 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..aebba48a83447 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,27 +331,17 @@ 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 (terms != null) { + if (termsEnum.seekExact(fieldName)) { + count += termsEnum.docFreq(); } } } 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 bc0d9ce2a84d7..b7d1beb4d1e06 100644 --- a/server/src/main/java/org/elasticsearch/index/shard/IndexShard.java +++ b/server/src/main/java/org/elasticsearch/index/shard/IndexShard.java @@ -1428,7 +1428,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..9794c123ca5eb 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 @@ -95,7 +95,7 @@ public FrozenEngine( fillSegmentStats(segmentReader, true, segmentsStats); } this.docsStats = docsStats(reader); - this.denseVectorStats = denseVectorStats(reader); + this.denseVectorStats = denseVectorStats(reader, null); this.sparseVectorStats = sparseVectorStats(reader, null); canMatchReader = ElasticsearchDirectoryReader.wrap( new RewriteCachingDirectoryReader(directory, reader.leaves(), null), @@ -334,7 +334,7 @@ public DocsStats docStats() { } @Override - public DenseVectorStats denseVectorStats() { + public DenseVectorStats denseVectorStats(MappingLookup mappingLookup) { return denseVectorStats; }