diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQuery.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQuery.java index a1dd1973b..988b46713 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQuery.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQuery.java @@ -38,14 +38,12 @@ public final class NeuralSparseQuery extends Query { */ @Override public String toString(String field) { - return "NeuralSparseQuery(" - + currentQuery.toString(field) - + "," - + highScoreTokenQuery.toString(field) - + ", " - + lowScoreTokenQuery.toString(field) - + ")"; - + return String.format( + "NeuralSparseQuery(%s,%s,%s)", + currentQuery.toString(field), + highScoreTokenQuery.toString(field), + lowScoreTokenQuery.toString(field) + ); } public Query rewrite(IndexSearcher indexSearcher) throws IOException { diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java index 2bd3a1d2f..e15fee7f7 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseQueryBuilder.java @@ -412,13 +412,14 @@ protected boolean doEquals(NeuralSparseQueryBuilder obj) { @Override protected int doHashCode() { - HashCodeBuilder builder = new HashCodeBuilder().append(fieldName).append(queryText).append(modelId).append(maxTokenScore); + HashCodeBuilder builder = new HashCodeBuilder().append(fieldName) + .append(queryText) + .append(modelId) + .append(maxTokenScore) + .append(neuralSparseTwoPhaseParameters); if (queryTokensSupplier != null) { builder.append(queryTokensSupplier.get()); } - if (Objects.nonNull(neuralSparseTwoPhaseParameters)) { - builder.append(neuralSparseTwoPhaseParameters.hashcode()); - } return builder.toHashCode(); } @@ -455,6 +456,9 @@ private Map getFilteredScoreTokens(Map queryTokens .stream() .filter(entry -> (aboveThreshold == (entry.getValue() >= threshold))) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + // This call will generate a filter from queryTokens. + // When aboveThreshold is true, will filter out all key-value pairs whose value >= threshold to a return map. + // When aboveThreshold is false, will filter out all key-value pairs whose value <= threshold to a return map. } private BooleanQuery buildFeatureFieldQueryFromTokens(Map tokens, String fieldName) { diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java index acd98e003..c645b0efd 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java @@ -63,7 +63,7 @@ public class NeuralSparseTwoPhaseParameters implements Writeable { * @param clusterService The opensearch clusterService. * @param settings The env settings to initialize. */ - public static void initialize(ClusterService clusterService, Settings settings) { + public static void initialize(final ClusterService clusterService, final Settings settings) { DEFAULT_ENABLED = NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_ENABLED.get(settings); DEFAULT_WINDOW_SIZE_EXPANSION = NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_WINDOW_SIZE_EXPANSION.get(settings); DEFAULT_PRUNING_RATIO = NeuralSearchSettings.NEURAL_SPARSE_TWO_PHASE_DEFAULT_PRUNING_RATIO.get(settings); @@ -96,14 +96,14 @@ public static NeuralSparseTwoPhaseParameters getDefaultSettings() { * @param in StreamInput to initialize object from * @throws IOException thrown if unable to read from input stream */ - public NeuralSparseTwoPhaseParameters(StreamInput in) throws IOException { + public NeuralSparseTwoPhaseParameters(final StreamInput in) throws IOException { window_size_expansion = in.readFloat(); pruning_ratio = in.readFloat(); enabled = in.readBoolean(); } @Override - public void writeTo(StreamOutput out) throws IOException { + public void writeTo(final StreamOutput out) throws IOException { out.writeFloat(window_size_expansion); out.writeFloat(pruning_ratio); out.writeBoolean(enabled); @@ -116,7 +116,7 @@ public void writeTo(StreamOutput out) throws IOException { * @return the given XContentBuilder with object content added. * @throws IOException if building the content fails. */ - public XContentBuilder doXContent(XContentBuilder builder) throws IOException { + public XContentBuilder doXContent(final XContentBuilder builder) throws IOException { builder.startObject(NAME.getPreferredName()); builder.field(WINDOW_SIZE_EXPANSION.getPreferredName(), window_size_expansion); builder.field(PRUNING_RATIO.getPreferredName(), pruning_ratio); @@ -132,7 +132,7 @@ public XContentBuilder doXContent(XContentBuilder builder) throws IOException { * @return a new instance of NeuralSparseTwoPhaseParameters initialized from the parser. * @throws IOException if parsing fails. */ - public static NeuralSparseTwoPhaseParameters parseFromXContent(XContentParser parser) throws IOException { + public static NeuralSparseTwoPhaseParameters parseFromXContent(final XContentParser parser) throws IOException { XContentParser.Token token; String currentFieldName = ""; NeuralSparseTwoPhaseParameters neuralSparseTwoPhaseParameters = NeuralSparseTwoPhaseParameters.getDefaultSettings(); @@ -177,7 +177,8 @@ public boolean equals(Object obj) { && Objects.equals(pruning_ratio, other.pruning_ratio); } - public int hashcode() { + @Override + public int hashCode() { HashCodeBuilder builder = new HashCodeBuilder().append(enabled).append(window_size_expansion).append(pruning_ratio); return builder.toHashCode(); } @@ -188,7 +189,7 @@ public int hashcode() { * @param neuralSparseTwoPhaseParameters The parameters to check. * @return true if enabled, false otherwise. */ - public static boolean isEnabled(NeuralSparseTwoPhaseParameters neuralSparseTwoPhaseParameters) { + public static boolean isEnabled(final NeuralSparseTwoPhaseParameters neuralSparseTwoPhaseParameters) { if (!isClusterOnOrAfterMinReqVersionForTwoPhaseSearchSupport() || Objects.isNull(neuralSparseTwoPhaseParameters)) { return false; } diff --git a/src/main/java/org/opensearch/neuralsearch/search/util/NeuralSparseTwoPhaseUtil.java b/src/main/java/org/opensearch/neuralsearch/search/util/NeuralSparseTwoPhaseUtil.java index f554e7194..49af06b5d 100644 --- a/src/main/java/org/opensearch/neuralsearch/search/util/NeuralSparseTwoPhaseUtil.java +++ b/src/main/java/org/opensearch/neuralsearch/search/util/NeuralSparseTwoPhaseUtil.java @@ -26,17 +26,22 @@ */ public class NeuralSparseTwoPhaseUtil { /** + * This function determine any neuralSparseQuery from query, extract lowScoreTokenQuery from each of them, + * And as that, the neuralSparseQuery be extracted 's currentQuery will be changed from allScoreTokenQuery to highScoreTokenQuery. + * Then build a QueryRescoreContext of these extra lowScoreTokenQuery and add the built QueryRescoreContext to the searchContext. + * Finally, the score of TopDocs will be sum of highScoreTokenQuery and lowScoreTokenQuery, which equals to allTokenQuery. * @param query The whole query include neuralSparseQuery to executed. * @param searchContext The searchContext with this query. */ - public static void addRescoreContextFromNeuralSparseQuery(final Query query, SearchContext searchContext) { + public static void addRescoreContextFromNeuralSparseQuery(final Query query, final SearchContext searchContext) { Map query2weight = new HashMap<>(); float windowSizeExpansion = populateQueryWeightsMapAndGetWindowSizeExpansion(query, query2weight, 1.0f, 1.0f); Query twoPhaseQuery; if (query2weight.isEmpty()) { return; - } else if (query2weight.size() == 1) { - Map.Entry entry = query2weight.entrySet().stream().findFirst().get(); + } + if (query2weight.size() == 1) { + Map.Entry entry = query2weight.entrySet().iterator().next(); twoPhaseQuery = new BoostQuery(entry.getKey(), entry.getValue()); } else { twoPhaseQuery = getNestedTwoPhaseQuery(query2weight); @@ -44,13 +49,13 @@ public static void addRescoreContextFromNeuralSparseQuery(final Query query, Sea int curWindowSize = (int) (searchContext.size() * windowSizeExpansion); if (curWindowSize < 0 || curWindowSize > NeuralSparseTwoPhaseParameters.MAX_WINDOW_SIZE) { throw new IllegalArgumentException( - "Two phase final windowSize " - + curWindowSize - + " out of score with limit " - + NeuralSparseTwoPhaseParameters.MAX_WINDOW_SIZE - + "." - + "You can change the value of cluster setting " - + "[plugins.neural_search.neural_sparse.two_phase.max_window_size] to a integer at least 50." + String.format( + "Two phase final windowSize %d out of score with limit %d. " + + "You can change the value of cluster setting [plugins.neural_search.neural_sparse.two_phase.max_window_size] " + + "to a integer at least 50.", + curWindowSize, + NeuralSparseTwoPhaseParameters.MAX_WINDOW_SIZE + ) ); } QueryRescorer.QueryRescoreContext rescoreContext = new QueryRescorer.QueryRescoreContext(curWindowSize); diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java index 01b10bf9a..106ced51e 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java @@ -47,9 +47,18 @@ public class NeuralSparseQueryIT extends BaseNeuralSearchIT { private static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field"; private static final String TEST_NEURAL_SPARSE_FIELD_NAME_NESTED = "nested.neural_sparse.field"; private static final List TEST_TOKENS = List.of("hello", "world", "a", "b", "c"); - + private static final List TWO_PHASE_TEST_TOKEN = List.of("hello", "world"); private static final Float DELTA = 1e-5f; private final Map testRankFeaturesDoc = TestUtils.createRandomTokenWeightMap(TEST_TOKENS); + private static final Map testFixedQueryTokens = new HashMap<>(); + private static final Supplier> testFixedQueryTokenSupplier = () -> testFixedQueryTokens; + static { + testFixedQueryTokens.put("hello", 5.0f); + testFixedQueryTokens.put("world", 4.0f); + testFixedQueryTokens.put("a", 3.0f); + testFixedQueryTokens.put("b", 2.0f); + testFixedQueryTokens.put("c", 1.0f); + } @Before public void setUp() throws Exception { @@ -304,18 +313,17 @@ public void testBooleanQuery_withMultipleSparseEncodingQueries() { */ @SneakyThrows public void testBooleanQuery_withMultipleSparseEncodingQueries_whenTwoPhaseEnabled() { - String modelId = null; try { initializeIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); - modelId = prepareSparseEncodingModel(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + Map randomTokenWeight = createRandomTokenWeightMap(TWO_PHASE_TEST_TOKEN); + Supplier> randomTokenWeightSupplier = () -> randomTokenWeight; NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(randomTokenWeightSupplier) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 2.0f, 0.4f)); NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_2) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(randomTokenWeightSupplier) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 2.0f, 0.4f)); boolQueryBuilder.should(sparseEncodingQueryBuilder1).should(sparseEncodingQueryBuilder2); @@ -324,10 +332,10 @@ public void testBooleanQuery_withMultipleSparseEncodingQueries_whenTwoPhaseEnabl Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, randomTokenWeightSupplier.get()); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME, null, null, null); } } @@ -489,35 +497,31 @@ protected void initializeIndexIfNotExist(String indexName) { */ @SneakyThrows public void testBasicQueryUsingQueryText_whenTwoPhaseEnabled_thenGetExpectedScore() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 2.0f, 0.4f)); Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @SneakyThrows public void testBasicQueryUsingQueryText_whenTwoPhaseEnabledAndDisabled_thenGetSameScore() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f); Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); @@ -526,7 +530,7 @@ public void testBasicQueryUsingQueryText_whenTwoPhaseEnabledAndDisabled_thenGetS sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters( new NeuralSparseTwoPhaseParameters().enabled(true).pruning_ratio(0.3f).window_size_expansion(6.0f) @@ -537,7 +541,7 @@ public void testBasicQueryUsingQueryText_whenTwoPhaseEnabledAndDisabled_thenGetS float scoreWithTwoPhase = objectToFloat(firstInnerHit.get("_score")); assertEquals(scoreWithTwoPhase, scoreWithoutTwoPhase, DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -550,14 +554,12 @@ public void testUpdateTwoPhaseSettings_whenTwoPhasedSettingsOverEdge_thenFail() @SneakyThrows public void testBasicQueryUsingQueryText_whenTwoPhaseParameterOverEdge_thenFail() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); // windows_size_expansion over edge [0.f NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(NeuralSparseTwoPhaseParameters.getDefaultSettings().window_size_expansion(-0.001f)); NeuralSparseQueryBuilder finalSparseEncodingQueryBuilder = sparseEncodingQueryBuilder; @@ -565,7 +567,7 @@ public void testBasicQueryUsingQueryText_whenTwoPhaseParameterOverEdge_thenFail( // pruning_ratio over edge [0.f sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(NeuralSparseTwoPhaseParameters.getDefaultSettings().pruning_ratio(-0.001f)); NeuralSparseQueryBuilder finalSparseEncodingQueryBuilder1 = sparseEncodingQueryBuilder; @@ -573,13 +575,13 @@ public void testBasicQueryUsingQueryText_whenTwoPhaseParameterOverEdge_thenFail( // pruning_ratio over edge 1.f] sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(NeuralSparseTwoPhaseParameters.getDefaultSettings().pruning_ratio(1.001f)); NeuralSparseQueryBuilder finalSparseEncodingQueryBuilder2 = sparseEncodingQueryBuilder; expectThrows(ResponseException.class, () -> search(TEST_BASIC_INDEX_NAME, finalSparseEncodingQueryBuilder2, 1)); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -614,23 +616,21 @@ public void testBasicQueryUsingQueryText_whenTwoPhaseParameterOverEdge_thenFail( */ @SneakyThrows public void testNeuralSparseQueryAsRescoreQuery_whenTwoPhase_thenGetExpectedScore() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 4.0f, 0.5f)); QueryBuilder queryBuilder = new MatchAllQueryBuilder(); Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, queryBuilder, sparseEncodingQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -673,14 +673,12 @@ public void testNeuralSparseQueryAsRescoreQuery_whenTwoPhase_thenGetExpectedScor */ @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhase_thenGetExpectedScore() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 4.0f, 0.2f)); boolQueryBuilder.should(sparseEncodingQueryBuilder); @@ -688,10 +686,10 @@ public void testMultiNeuralSparseQuery_whenTwoPhase_thenGetExpectedScore() { Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 4 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 4 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -730,14 +728,12 @@ public void testNeuralSparseQuery_whenDifferentTwoPhaseParameter_thenGetDifferen @SneakyThrows public void testMultiNeuralSparseQuery_whenTwoPhaseAndFilter_thenGetExpectedScore() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(2.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 5.0f, 0.8f)); boolQueryBuilder.should(sparseEncodingQueryBuilder); @@ -745,10 +741,10 @@ public void testMultiNeuralSparseQuery_whenTwoPhaseAndFilter_thenGetExpectedScor Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -833,13 +829,11 @@ public void testMultiNeuralSparseQuery_whenTwoPhaseAndNoLowScoreToken_thenGetExp */ @SneakyThrows public void testNeuralSParseQuery_whenTwoPhaseAndNestedInConstantScoreQuery_thenSuccess() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(1.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 5.0f, 0.6f)); ConstantScoreQueryBuilder constantScoreQueryBuilder = new ConstantScoreQueryBuilder(sparseEncodingQueryBuilder); @@ -848,7 +842,7 @@ public void testNeuralSParseQuery_whenTwoPhaseAndNestedInConstantScoreQuery_then assertEquals("1", firstInnerHit.get("_id")); assertEquals(1.0f, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -876,13 +870,11 @@ public void testNeuralSParseQuery_whenTwoPhaseAndNestedInConstantScoreQuery_then */ @SneakyThrows public void testNeuralSParseQuery_whenTwoPhaseAndNestedInDisjunctionMaxQuery_thenSuccess() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(5.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 5.0f, 0.6f)); DisMaxQueryBuilder disMaxQueryBuilder = new DisMaxQueryBuilder(); @@ -891,10 +883,10 @@ public void testNeuralSParseQuery_whenTwoPhaseAndNestedInDisjunctionMaxQuery_the Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, disMaxQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 5f * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 5f * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } } @@ -917,13 +909,11 @@ public void testNeuralSParseQuery_whenTwoPhaseAndNestedInDisjunctionMaxQuery_the */ @SneakyThrows public void testNeuralSParseQuery_whenTwoPhaseAndNestedInFunctionScoreQuery_thenSuccess() { - String modelId = null; try { initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); - modelId = prepareSparseEncodingModel(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) .queryText(TEST_QUERY_TEXT) - .modelId(modelId) + .queryTokensSupplier(testFixedQueryTokenSupplier) .boost(5.0f) .neuralSparseTwoPhaseParameters(getCustomTwoPhaseParameter(true, 5.0f, 0.6f)); FunctionScoreQueryBuilder functionScoreQueryBuilder = new FunctionScoreQueryBuilder(sparseEncodingQueryBuilder); @@ -931,10 +921,10 @@ public void testNeuralSParseQuery_whenTwoPhaseAndNestedInFunctionScoreQuery_then Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, functionScoreQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 10f * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 10f * computeExpectedScore(testRankFeaturesDoc, testFixedQueryTokens); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } finally { - wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, modelId, null); + wipeOfTestResources(TEST_BASIC_INDEX_NAME, null, null, null); } }