From 4f29bd85400529a6011637557a33d8ff2eec1bdf Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Sat, 30 Nov 2024 17:21:11 +0000 Subject: [PATCH] iter --- .../index/mapper/DocumentParserContext.java | 30 ++++++++-------- .../highlight/DefaultHighlighter.java | 2 +- .../index/mapper/MapperTestCase.java | 2 -- x-pack/plugin/core/build.gradle | 2 +- .../xpack/core/XPackClientPlugin.java | 23 ++++++++++++ ...ery.java => SparseVectorQueryWrapper.java} | 8 ++--- .../core/ml/search/WeightedTokensUtils.java | 4 +-- .../search/SparseVectorQueryBuilderTests.java | 2 +- .../WeightedTokensQueryBuilderTests.java | 8 ++--- .../inference/src/main/java/module-info.java | 1 - .../xpack/inference/InferencePlugin.java | 20 +---------- .../xpack/inference/chunking/Chunker.java | 4 --- .../chunking/SentenceBoundaryChunker.java | 35 +++---------------- .../chunking/WordBoundaryChunker.java | 5 --- .../highlight/SemanticTextHighlighter.java | 4 +-- ...ffsetField.java => OffsetSourceField.java} | 4 +-- .../mapper/OffsetSourceFieldMapper.java | 2 +- .../queries/SemanticQueryBuilderTests.java | 6 ++-- .../SparseVectorQueryBuilderTests.java | 4 +-- .../WeightedTokensQueryBuilderTests.java | 10 +++--- 20 files changed, 71 insertions(+), 105 deletions(-) rename x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/{SparseVectorQuery.java => SparseVectorQueryWrapper.java} (87%) rename x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/{OffsetField.java => OffsetSourceField.java} (94%) diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java b/server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java index 9f4dbabf4e700..691402cf682ca 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java +++ b/server/src/main/java/org/elasticsearch/index/mapper/DocumentParserContext.java @@ -349,21 +349,6 @@ public final DocumentParserContext addIgnoredFieldFromContext(IgnoredSourceField return this; } - /** - * Called by {@link InferenceMetadataFieldsMapper} to indicate whether the metadata field is present - * in _source. - */ - public void markInferenceMetadataField() { - this.hasInferenceMetadata = true; - } - - /** - * Returns whether the _source contains an inference metadata field. - */ - public final boolean hasInferenceMetadataField() { - return hasInferenceMetadata; - } - /** * Wraps {@link XContentDataHelper#encodeToken}, disabling dot expansion from {@link DotExpandingXContentParser}. * This helps avoid producing duplicate names in the same scope, due to expanding dots to objects. @@ -679,6 +664,21 @@ public boolean isWithinInferenceMetadata() { return false; } + /** + * Called by {@link InferenceMetadataFieldsMapper} to indicate whether the metadata field is present + * in _source. + */ + public void markInferenceMetadataField() { + this.hasInferenceMetadata = true; + } + + /** + * Returns whether the _source contains an inference metadata field. + */ + public final boolean hasInferenceMetadataField() { + return hasInferenceMetadata; + } + boolean inArrayScope() { return currentScope == Scope.ARRAY; } diff --git a/server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/DefaultHighlighter.java b/server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/DefaultHighlighter.java index e98ad8a5bc572..75f8e5588761a 100644 --- a/server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/DefaultHighlighter.java +++ b/server/src/main/java/org/elasticsearch/search/fetch/subphase/highlight/DefaultHighlighter.java @@ -235,7 +235,7 @@ protected static String convertFieldValue(MappedFieldType type, Object value) { } } - public static String mergeFieldValues(List fieldValues, char valuesSeparator) { + protected static String mergeFieldValues(List fieldValues, char valuesSeparator) { // postings highlighter accepts all values in a single string, as offsets etc. need to match with content // loaded from stored fields, we merge all values using a proper separator String rawValue = Strings.collectionToDelimitedString(fieldValues, String.valueOf(valuesSeparator)); diff --git a/test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java b/test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java index 02aa483c4e214..29bb3b15a9f86 100644 --- a/test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java +++ b/test/framework/src/main/java/org/elasticsearch/index/mapper/MapperTestCase.java @@ -1166,10 +1166,8 @@ public void testSupportsParsingObject() throws IOException { Object sampleValueForDocument = getSampleObjectForDocument(); assertThat(sampleValueForDocument, instanceOf(Map.class)); SourceToParse source = source(builder -> { - builder.startObject(InferenceMetadataFieldsMapper.NAME); builder.field("field"); builder.value(sampleValueForDocument); - builder.endObject(); }); ParsedDocument doc = mapper.parse(source); assertNotNull(doc); diff --git a/x-pack/plugin/core/build.gradle b/x-pack/plugin/core/build.gradle index 5d2c8c9811c0a..51d770936e64e 100644 --- a/x-pack/plugin/core/build.gradle +++ b/x-pack/plugin/core/build.gradle @@ -52,6 +52,7 @@ dependencies { api "commons-codec:commons-codec:${versions.commonscodec}" testImplementation project(path: ':modules:aggregations') testImplementation project(path: ':modules:data-streams') + testImplementation project(':modules:mapper-extras') // security deps api 'com.unboundid:unboundid-ldapsdk:6.0.3' @@ -68,7 +69,6 @@ dependencies { testImplementation project(path: ':modules:analysis-common') testImplementation project(path: ':modules:rest-root') testImplementation project(path: ':modules:health-shards-availability') - testImplementation project(path: ':modules:mapper-extras') // Needed for Fips140ProviderVerificationTests testCompileOnly('org.bouncycastle:bc-fips:1.0.2.5') diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackClientPlugin.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackClientPlugin.java index b84c1b6a48331..23f9e91dc32da 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackClientPlugin.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/XPackClientPlugin.java @@ -11,6 +11,7 @@ import org.elasticsearch.cluster.metadata.Metadata; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.settings.Setting; +import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.license.LicensesMetadata; import org.elasticsearch.persistent.PersistentTaskParams; import org.elasticsearch.persistent.PersistentTaskState; @@ -70,6 +71,9 @@ import org.elasticsearch.xpack.core.ml.job.config.JobTaskState; import org.elasticsearch.xpack.core.ml.job.snapshot.upgrade.SnapshotUpgradeTaskParams; import org.elasticsearch.xpack.core.ml.job.snapshot.upgrade.SnapshotUpgradeTaskState; +import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryBuilder; +import org.elasticsearch.xpack.core.ml.search.TextExpansionQueryBuilder; +import org.elasticsearch.xpack.core.ml.search.WeightedTokensQueryBuilder; import org.elasticsearch.xpack.core.monitoring.MonitoringFeatureSetUsage; import org.elasticsearch.xpack.core.rollup.RollupFeatureSetUsage; import org.elasticsearch.xpack.core.rollup.RollupField; @@ -390,4 +394,23 @@ public List getNamedXContent() { ) ); } + + @Override + public List> getQueries() { + return List.of( + new QuerySpec<>(SparseVectorQueryBuilder.NAME, SparseVectorQueryBuilder::new, SparseVectorQueryBuilder::fromXContent), + new QuerySpec( + TextExpansionQueryBuilder.NAME, + TextExpansionQueryBuilder::new, + TextExpansionQueryBuilder::fromXContent + ), + // TODO: The WeightedTokensBuilder is slated for removal after the SparseVectorQueryBuilder is available. + // The logic to create a Boolean query based on weighted tokens will remain and/or be moved to server. + new SearchPlugin.QuerySpec( + WeightedTokensQueryBuilder.NAME, + WeightedTokensQueryBuilder::new, + WeightedTokensQueryBuilder::fromXContent + ) + ); + } } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQuery.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryWrapper.java similarity index 87% rename from x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQuery.java rename to x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryWrapper.java index 2cb4d6777dcb9..3a3940524a8d4 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQuery.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryWrapper.java @@ -17,11 +17,11 @@ import java.io.IOException; import java.util.Objects; -public class SparseVectorQuery extends Query { +public class SparseVectorQueryWrapper extends Query { private final String fieldName; private final Query termsQuery; - public SparseVectorQuery(String fieldName, Query termsQuery) { + public SparseVectorQueryWrapper(String fieldName, Query termsQuery) { this.fieldName = fieldName; this.termsQuery = termsQuery; } @@ -34,7 +34,7 @@ public Query getTermsQuery() { public Query rewrite(IndexSearcher indexSearcher) throws IOException { var rewrite = termsQuery.rewrite(indexSearcher); if (rewrite != termsQuery) { - return new SparseVectorQuery(fieldName, rewrite); + return new SparseVectorQueryWrapper(fieldName, rewrite); } return this; } @@ -61,7 +61,7 @@ public boolean equals(Object obj) { if (sameClassAs(obj) == false) { return false; } - SparseVectorQuery that = (SparseVectorQuery) obj; + SparseVectorQueryWrapper that = (SparseVectorQueryWrapper) obj; return fieldName.equals(that.fieldName) && termsQuery.equals(that.termsQuery); } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensUtils.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensUtils.java index 0fcd07ed8ce08..1c2ac23151e6e 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensUtils.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensUtils.java @@ -35,7 +35,7 @@ public static Query queryBuilderWithAllTokens( for (var token : tokens) { qb.add(new BoostQuery(ft.termQuery(token.token(), context), token.weight()), BooleanClause.Occur.SHOULD); } - return new SparseVectorQuery(fieldName, qb.setMinimumNumberShouldMatch(1).build()); + return new SparseVectorQueryWrapper(fieldName, qb.setMinimumNumberShouldMatch(1).build()); } public static Query queryBuilderWithPrunedTokens( @@ -69,7 +69,7 @@ public static Query queryBuilderWithPrunedTokens( } } - return new SparseVectorQuery(fieldName, qb.setMinimumNumberShouldMatch(1).build()); + return new SparseVectorQueryWrapper(fieldName, qb.setMinimumNumberShouldMatch(1).build()); } /** diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilderTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilderTests.java index f88e7467c29b2..a6105d5994ab5 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilderTests.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/SparseVectorQueryBuilderTests.java @@ -234,7 +234,7 @@ private void testDoToQuery(SparseVectorQueryBuilder queryBuilder, SearchExecutio // It's possible that all documents were pruned for aggressive pruning configurations assertTrue(query instanceof BooleanQuery || query instanceof MatchNoDocsQuery); } else { - assertTrue(query instanceof SparseVectorQuery); + assertTrue(query instanceof SparseVectorQueryWrapper); } } diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensQueryBuilderTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensQueryBuilderTests.java index 7372def52355e..6f61749f35fb3 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensQueryBuilderTests.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/WeightedTokensQueryBuilderTests.java @@ -271,8 +271,8 @@ public void testPruningIsAppliedCorrectly() throws IOException { } private void assertCorrectLuceneQuery(String name, Query query, List expectedFeatureFields) { - assertTrue(query instanceof SparseVectorQuery); - Query termsQuery = ((SparseVectorQuery) query).getTermsQuery(); + assertTrue(query instanceof SparseVectorQueryWrapper); + Query termsQuery = ((SparseVectorQueryWrapper) query).getTermsQuery(); assertTrue(termsQuery instanceof BooleanQuery); List booleanClauses = ((BooleanQuery) termsQuery).clauses(); assertEquals( @@ -345,8 +345,8 @@ public void testMustRewrite() throws IOException { @Override protected void doAssertLuceneQuery(WeightedTokensQueryBuilder queryBuilder, Query query, SearchExecutionContext context) { - assertThat(query, instanceOf(SparseVectorQuery.class)); - Query termsQuery = ((SparseVectorQuery) query).getTermsQuery(); + assertThat(query, instanceOf(SparseVectorQueryWrapper.class)); + Query termsQuery = ((SparseVectorQueryWrapper) query).getTermsQuery(); assertThat(termsQuery, instanceOf(BooleanQuery.class)); BooleanQuery booleanQuery = (BooleanQuery) termsQuery; assertEquals(booleanQuery.getMinimumNumberShouldMatch(), 1); diff --git a/x-pack/plugin/inference/src/main/java/module-info.java b/x-pack/plugin/inference/src/main/java/module-info.java index 13f54d5d580bd..53974657e4e23 100644 --- a/x-pack/plugin/inference/src/main/java/module-info.java +++ b/x-pack/plugin/inference/src/main/java/module-info.java @@ -34,7 +34,6 @@ requires software.amazon.awssdk.retries.api; requires org.reactivestreams; requires org.elasticsearch.logging; - requires org.apache.lucene.highlighter; exports org.elasticsearch.xpack.inference.action; exports org.elasticsearch.xpack.inference.registry; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java index 8f1d201bba0e8..afeb50ed5e634 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java @@ -77,9 +77,6 @@ import org.elasticsearch.xpack.inference.mapper.OffsetSourceMetaFieldMapper; import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder; -import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryBuilder; -import org.elasticsearch.xpack.core.ml.search.TextExpansionQueryBuilder; -import org.elasticsearch.xpack.core.ml.search.WeightedTokensQueryBuilder; import org.elasticsearch.xpack.inference.rank.random.RandomRankBuilder; import org.elasticsearch.xpack.inference.rank.random.RandomRankRetrieverBuilder; import org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankBuilder; @@ -426,22 +423,7 @@ public Collection getMappedActionFilters() { } public List> getQueries() { - return List.of( - new QuerySpec<>(SemanticQueryBuilder.NAME, SemanticQueryBuilder::new, SemanticQueryBuilder::fromXContent), - new QuerySpec<>(SparseVectorQueryBuilder.NAME, SparseVectorQueryBuilder::new, SparseVectorQueryBuilder::fromXContent), - new QuerySpec( - TextExpansionQueryBuilder.NAME, - TextExpansionQueryBuilder::new, - TextExpansionQueryBuilder::fromXContent - ), - // TODO: The WeightedTokensBuilder is slated for removal after the SparseVectorQueryBuilder is available. - // The logic to create a Boolean query based on weighted tokens will remain and/or be moved to server. - new QuerySpec( - WeightedTokensQueryBuilder.NAME, - WeightedTokensQueryBuilder::new, - WeightedTokensQueryBuilder::fromXContent - ) - ); + return List.of(new QuerySpec<>(SemanticQueryBuilder.NAME, SemanticQueryBuilder::new, SemanticQueryBuilder::fromXContent)); } @Override diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/Chunker.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/Chunker.java index 01919ef19c6fd..af7c706c807ec 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/Chunker.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/Chunker.java @@ -12,9 +12,5 @@ import java.util.List; public interface Chunker { - record Chunk(int startOffset, int endOffset) {} - List chunk(String input, ChunkingSettings chunkingSettings); - - List chunkOffset(String input, ChunkingSettings chunkingSettings); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/SentenceBoundaryChunker.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/SentenceBoundaryChunker.java index c0315d0c56daf..5df940d6a3fba 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/SentenceBoundaryChunker.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/SentenceBoundaryChunker.java @@ -15,7 +15,6 @@ import java.util.ArrayList; import java.util.List; import java.util.Locale; -import java.util.stream.Collectors; /** * Split text into chunks aligned on sentence boundaries. @@ -59,20 +58,6 @@ public List chunk(String input, ChunkingSettings chunkingSettings) { } } - @Override - public List chunkOffset(String input, ChunkingSettings chunkingSettings) { - if (chunkingSettings instanceof SentenceBoundaryChunkingSettings sentenceBoundaryChunkingSettings) { - return chunkOffset(input, sentenceBoundaryChunkingSettings.maxChunkSize, sentenceBoundaryChunkingSettings.sentenceOverlap > 0); - } else { - throw new IllegalArgumentException( - Strings.format( - "SentenceBoundaryChunker can't use ChunkingSettings with strategy [%s]", - chunkingSettings.getChunkingStrategy() - ) - ); - } - } - /** * Break the input text into small chunks on sentence boundaries. * @@ -81,19 +66,7 @@ public List chunkOffset(String input, ChunkingSettings chunkingSettings) * @return The input text chunked */ public List chunk(String input, int maxNumberWordsPerChunk, boolean includePrecedingSentence) { - var chunks = chunkOffset(input, maxNumberWordsPerChunk, includePrecedingSentence); - return chunks.stream().map(c -> input.substring(c.startOffset(), c.endOffset())).collect(Collectors.toList()); - } - - /** - * Break the input text into small chunks on sentence boundaries. - * - * @param input Text to chunk - * @param maxNumberWordsPerChunk Maximum size of the chunk - * @return The input text chunked - */ - public List chunkOffset(String input, int maxNumberWordsPerChunk, boolean includePrecedingSentence) { - var chunks = new ArrayList(); + var chunks = new ArrayList(); sentenceIterator.setText(input); wordIterator.setText(input); @@ -118,7 +91,7 @@ public List chunkOffset(String input, int maxNumberWordsPerChunk, boolean int nextChunkWordCount = wordsInSentenceCount; if (chunkWordCount > 0) { // add a new chunk containing all the input up to this sentence - chunks.add(new Chunk(chunkStart, chunkEnd)); + chunks.add(input.substring(chunkStart, chunkEnd)); if (includePrecedingSentence) { if (wordsInPrecedingSentenceCount + wordsInSentenceCount > maxNumberWordsPerChunk) { @@ -154,7 +127,7 @@ public List chunkOffset(String input, int maxNumberWordsPerChunk, boolean for (; i < sentenceSplits.size() - 1; i++) { // Because the substring was passed to splitLongSentence() // the returned positions need to be offset by chunkStart - chunks.add(new Chunk(chunkStart + sentenceSplits.get(i).start(), chunkStart + sentenceSplits.get(i).end())); + chunks.add(input.substring(chunkStart + sentenceSplits.get(i).start(), chunkStart + sentenceSplits.get(i).end())); } // The final split is partially filled. // Set the next chunk start to the beginning of the @@ -178,7 +151,7 @@ public List chunkOffset(String input, int maxNumberWordsPerChunk, boolean } if (chunkWordCount > 0) { - chunks.add(new Chunk(chunkStart, input.length())); + chunks.add(input.substring(chunkStart)); } return chunks; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/WordBoundaryChunker.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/WordBoundaryChunker.java index 20ab2dbdd0fc9..c9c752b9aabbc 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/WordBoundaryChunker.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/WordBoundaryChunker.java @@ -55,11 +55,6 @@ public List chunk(String input, ChunkingSettings chunkingSettings) { } } - @Override - public List chunkOffset(String input, ChunkingSettings chunkingSettings) { - return List.of(); - } - /** * Break the input text into small chunks as dictated * by the chunking parameters diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java index 09723d76bb358..eda77a6269513 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/highlight/SemanticTextHighlighter.java @@ -35,7 +35,7 @@ import org.elasticsearch.xpack.inference.mapper.OffsetSourceMetaFieldMapper; import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; import org.elasticsearch.xpack.inference.mapper.SemanticTextUtils; -import org.elasticsearch.xpack.core.ml.search.SparseVectorQuery; +import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryWrapper; import java.io.IOException; import java.util.ArrayList; @@ -217,7 +217,7 @@ public void consumeTerms(Query query, Term... terms) { @Override public QueryVisitor getSubVisitor(BooleanClause.Occur occur, Query parent) { - if (parent instanceof SparseVectorQuery sparseVectorQuery) { + if (parent instanceof SparseVectorQueryWrapper sparseVectorQuery) { queries.add(sparseVectorQuery.getTermsQuery()); } return this; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetField.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceField.java similarity index 94% rename from x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetField.java rename to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceField.java index 3449a21b51104..4439d6a00880c 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetField.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceField.java @@ -15,7 +15,7 @@ import org.apache.lucene.document.FieldType; import org.apache.lucene.index.IndexOptions; -public final class OffsetField extends Field { +public final class OffsetSourceField extends Field { private static final FieldType FIELD_TYPE = new FieldType(); @@ -28,7 +28,7 @@ public final class OffsetField extends Field { private int startOffset; private int endOffset; - public OffsetField(String fieldName, String sourceFieldName, int startOffset, int endOffset) { + public OffsetSourceField(String fieldName, String sourceFieldName, int startOffset, int endOffset) { super(fieldName, sourceFieldName, FIELD_TYPE); this.startOffset = startOffset; this.endOffset = endOffset; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceFieldMapper.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceFieldMapper.java index bd5be953b69a8..7d36505bb2209 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceFieldMapper.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/OffsetSourceFieldMapper.java @@ -151,7 +151,7 @@ protected void parseCreateField(DocumentParserContext context) throws IOExceptio throw new IllegalArgumentException("Unkown field name [" + fieldName + "]"); } } - context.doc().addWithKey(fullPath(), new OffsetField(NAME, fullPath() + "." + sourceFieldName, startOffset, endOffset)); + context.doc().addWithKey(fullPath(), new OffsetSourceField(NAME, fullPath() + "." + sourceFieldName, startOffset, endOffset)); } finally { context.path().setWithinLeafObject(isWithinLeafObject); } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java index 2015d1c938c56..7e6d0775049d4 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java @@ -52,7 +52,7 @@ import org.elasticsearch.xpack.core.ml.inference.MlInferenceNamedXContentProvider; import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults; import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults; -import org.elasticsearch.xpack.core.ml.search.SparseVectorQuery; +import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryWrapper; import org.elasticsearch.xpack.core.ml.search.WeightedToken; import org.elasticsearch.xpack.inference.InferencePlugin; import org.elasticsearch.xpack.inference.mapper.SemanticTextField; @@ -190,8 +190,8 @@ protected void doAssertLuceneQuery(SemanticQueryBuilder queryBuilder, Query quer } private void assertSparseEmbeddingLuceneQuery(Query query) { - assertThat(query, instanceOf(SparseVectorQuery.class)); - Query termsQuery = ((SparseVectorQuery) query).getTermsQuery(); + assertThat(query, instanceOf(SparseVectorQueryWrapper.class)); + Query termsQuery = ((SparseVectorQueryWrapper) query).getTermsQuery(); Query innerQuery = assertOuterBooleanQuery(termsQuery); assertThat(innerQuery, instanceOf(BooleanQuery.class)); BooleanQuery innerBooleanQuery = (BooleanQuery) innerQuery; diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SparseVectorQueryBuilderTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SparseVectorQueryBuilderTests.java index 832514929e42c..24514ab272232 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SparseVectorQueryBuilderTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SparseVectorQueryBuilderTests.java @@ -40,7 +40,7 @@ import org.elasticsearch.xpack.core.ml.action.InferModelAction; import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings; import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults; -import org.elasticsearch.xpack.core.ml.search.SparseVectorQuery; +import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryWrapper; import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryBuilder; import org.elasticsearch.xpack.core.ml.search.TokenPruningConfig; import org.elasticsearch.xpack.core.ml.search.WeightedToken; @@ -239,7 +239,7 @@ private void testDoToQuery(SparseVectorQueryBuilder queryBuilder, SearchExecutio // It's possible that all documents were pruned for aggressive pruning configurations assertTrue(query instanceof BooleanQuery || query instanceof MatchNoDocsQuery); } else { - assertTrue(query instanceof SparseVectorQuery); + assertTrue(query instanceof SparseVectorQueryWrapper); } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/WeightedTokensQueryBuilderTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/WeightedTokensQueryBuilderTests.java index 20cdf7c27dc3d..77e0324dec92b 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/WeightedTokensQueryBuilderTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/WeightedTokensQueryBuilderTests.java @@ -35,7 +35,7 @@ import org.elasticsearch.xpack.core.ml.action.InferModelAction; import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings; import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults; -import org.elasticsearch.xpack.core.ml.search.SparseVectorQuery; +import org.elasticsearch.xpack.core.ml.search.SparseVectorQueryWrapper; import org.elasticsearch.xpack.core.ml.search.TokenPruningConfig; import org.elasticsearch.xpack.core.ml.search.WeightedToken; import org.elasticsearch.xpack.core.ml.search.WeightedTokensQueryBuilder; @@ -276,8 +276,8 @@ public void testPruningIsAppliedCorrectly() throws IOException { } private void assertCorrectLuceneQuery(String name, Query query, List expectedFeatureFields) { - assertTrue(query instanceof SparseVectorQuery); - Query termsQuery = ((SparseVectorQuery) query).getTermsQuery(); + assertTrue(query instanceof SparseVectorQueryWrapper); + Query termsQuery = ((SparseVectorQueryWrapper) query).getTermsQuery(); assertTrue(termsQuery instanceof BooleanQuery); List booleanClauses = ((BooleanQuery) termsQuery).clauses(); assertEquals( @@ -350,8 +350,8 @@ public void testMustRewrite() throws IOException { @Override protected void doAssertLuceneQuery(WeightedTokensQueryBuilder queryBuilder, Query query, SearchExecutionContext context) { - assertThat(query, instanceOf(SparseVectorQuery.class)); - Query termsQuery = ((SparseVectorQuery) query).getTermsQuery(); + assertThat(query, instanceOf(SparseVectorQueryWrapper.class)); + Query termsQuery = ((SparseVectorQueryWrapper) query).getTermsQuery(); assertThat(termsQuery, instanceOf(BooleanQuery.class)); BooleanQuery booleanQuery = (BooleanQuery) termsQuery; assertEquals(booleanQuery.getMinimumNumberShouldMatch(), 1);