From bf264edb70fbfbb4c66bbbc4c2acab264c4e0071 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Wed, 29 Nov 2023 03:10:29 -0800 Subject: [PATCH] Formatting of Test folder and Initial BWC Test Commit Signed-off-by: Varun Jain --- build.gradle | 10 + qa/build.gradle | 6 +- .../AbstractRestartUpgradeRestTestCase.java | 62 ++++ .../neuralsearch/bwc/TextSearchIT.java | 67 ++++ .../bwc/AbstractRollingUpgradeTestCase.java | 90 ++++++ repositories.gradle | 11 + .../plugin/NeuralSearchTests.java | 4 +- .../NeuralQueryEnricherProcessorIT.java | 45 +-- .../processor/NormalizationProcessorIT.java | 157 +-------- .../processor/ScoreCombinationIT.java | 153 +-------- .../processor/ScoreNormalizationIT.java | 97 +----- .../processor/SparseEncodingProcessIT.java | 2 +- .../processor/TextEmbeddingProcessorIT.java | 21 +- .../TextImageEmbeddingProcessorIT.java | 23 +- .../neuralsearch/query/HybridQueryIT.java | 124 +------- .../neuralsearch/query/NeuralQueryIT.java | 171 +++------- .../query/NeuralSparseQueryIT.java | 112 ++----- .../neuralsearch}/BaseNeuralSearchIT.java | 297 +++++++++++++++++- .../neuralsearch}/BaseSparseEncodingIT.java | 52 ++- .../OpenSearchSecureRestTestCase.java | 0 .../opensearch/neuralsearch/TestUtils.java | 70 ++++- 21 files changed, 804 insertions(+), 770 deletions(-) create mode 100644 qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java create mode 100644 qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java create mode 100644 repositories.gradle rename src/{test/java/org/opensearch/neuralsearch/common => testFixtures/java/org/opensearch/neuralsearch}/BaseNeuralSearchIT.java (69%) rename src/{test/java/org/opensearch/neuralsearch/common => testFixtures/java/org/opensearch/neuralsearch}/BaseSparseEncodingIT.java (67%) rename src/{test => testFixtures}/java/org/opensearch/neuralsearch/OpenSearchSecureRestTestCase.java (100%) rename src/{test => testFixtures}/java/org/opensearch/neuralsearch/TestUtils.java (70%) diff --git a/build.gradle b/build.gradle index 7220fda23..d7eaae444 100644 --- a/build.gradle +++ b/build.gradle @@ -11,6 +11,7 @@ import org.opensearch.gradle.test.RestIntegTestTask import java.util.concurrent.Callable apply plugin: 'java' +apply plugin: 'java-test-fixtures' apply plugin: 'idea' apply plugin: 'opensearch.opensearchplugin' apply plugin: 'opensearch.pluginzip' @@ -123,6 +124,7 @@ allprojects { version = "${opensearch_build}" targetCompatibility = JavaVersion.VERSION_11 sourceCompatibility = JavaVersion.VERSION_11 + apply from: rootProject.file('repositories.gradle').absoluteFile } repositories { @@ -145,6 +147,7 @@ dependencies { compileOnly fileTree(dir: knnJarDirectory, include: '*.jar') api group: 'org.opensearch', name:'opensearch-ml-client', version: "${opensearch_build}" implementation group: 'org.apache.commons', name: 'commons-lang3', version: '3.10' + testFixturesImplementation "org.opensearch.test:framework:${opensearch_version}" // ml-common excluded reflection for runtime so we need to add it by ourselves. // https://github.com/opensearch-project/ml-commons/commit/464bfe34c66d7a729a00dd457f03587ea4e504d9 // TODO: Remove following three lines of dependencies if ml-common include them in their jar @@ -154,6 +157,10 @@ dependencies { runtimeOnly group: 'org.apache.commons', name: 'commons-text', version: '1.10.0' runtimeOnly group: 'com.google.code.gson', name: 'gson', version: '2.10.1' runtimeOnly group: 'org.json', name: 'json', version: '20231013' + testFixturesImplementation "org.opensearch:common-utils:${version}" + testFixturesImplementation group: 'org.apache.commons', name: 'commons-lang3', version: '3.10' + testFixturesCompileOnly group: 'com.google.guava', name: 'guava', version:'32.0.1-jre' + testFixturesCompileOnly fileTree(dir: knnJarDirectory, include: '*.jar') } // In order to add the jar to the classpath, we need to unzip the @@ -173,6 +180,9 @@ compileJava { compileTestJava { options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor']) } +compileTestFixturesJava { + options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor']) +} def opensearch_tmp_dir = rootProject.file('build/private/opensearch_tmp').absoluteFile opensearch_tmp_dir.mkdirs() diff --git a/qa/build.gradle b/qa/build.gradle index ca8f7cddc..df1273df6 100644 --- a/qa/build.gradle +++ b/qa/build.gradle @@ -9,6 +9,7 @@ import java.nio.file.Files apply plugin: 'opensearch.testclusters' apply plugin: 'opensearch.build' apply plugin: 'opensearch.rest-test' +apply plugin: 'io.freefair.lombok' // Disable a few tasks that come with build build.enabled = false @@ -19,11 +20,10 @@ dependenciesInfo.enabled = false dependencies { api "org.opensearch:opensearch:${opensearch_version}" - //compileOnly "org.opensearch.plugin:opensearch-scripting-painless-spi:${versions.opensearch}" - //api group: 'commons-lang', name: 'commons-lang', version: '2.6' - + api group: 'commons-lang', name: 'commons-lang', version: '2.6' api "org.apache.logging.log4j:log4j-api:${versions.log4j}" api "org.apache.logging.log4j:log4j-core:${versions.log4j}" + api "junit:junit:${versions.junit}" testImplementation "org.opensearch.test:framework:${opensearch_version}" testImplementation(testFixtures(rootProject)) diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java new file mode 100644 index 000000000..cc1d66084 --- /dev/null +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java @@ -0,0 +1,62 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.bwc; + +import java.util.Locale; +import java.util.Optional; +import org.junit.Before; +import org.opensearch.common.settings.Settings; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; +import static org.opensearch.neuralsearch.TestUtils.CLIENT_TIMEOUT_VALUE; +import static org.opensearch.neuralsearch.TestUtils.RESTART_UPGRADE_OLD_CLUSTER; +import static org.opensearch.neuralsearch.TestUtils.BWC_VERSION; +import static org.opensearch.neuralsearch.TestUtils.NEURAL_SEARCH_BWC_PREFIX; +import org.opensearch.test.rest.OpenSearchRestTestCase; + +public abstract class AbstractRestartUpgradeRestTestCase extends BaseNeuralSearchIT { + protected static String testIndex; + + @Before + protected void setIndex() { + // Creating index name by concatenating "knn-bwc-" prefix with test method name + // for all the tests in this sub-project + testIndex = NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); + } + + @Override + protected final boolean preserveIndicesUponCompletion() { + return true; + } + + @Override + protected final boolean preserveReposUponCompletion() { + return true; + } + + @Override + protected boolean preserveTemplatesUponCompletion() { + return true; + } + + @Override + protected final Settings restClientSettings() { + return Settings.builder() + .put(super.restClientSettings()) + // increase the timeout here to 90 seconds to handle long waits for a green + // cluster health. the waits for green need to be longer than a minute to + // account for delayed shards + .put(OpenSearchRestTestCase.CLIENT_SOCKET_TIMEOUT, CLIENT_TIMEOUT_VALUE) + .build(); + } + + protected static final boolean isRunningAgainstOldCluster() { + return Boolean.parseBoolean(System.getProperty(RESTART_UPGRADE_OLD_CLUSTER)); + } + + protected final Optional getBWCVersion() { + return Optional.ofNullable(System.getProperty(BWC_VERSION, null)); + } +} diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/TextSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/TextSearchIT.java index e69de29bb..baac0ef49 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/TextSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/TextSearchIT.java @@ -0,0 +1,67 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.bwc; + +import java.util.Map; +import lombok.SneakyThrows; +import static org.opensearch.neuralsearch.TestUtils.*; +import static org.opensearch.neuralsearch.TestUtils.objectToFloat; +import org.opensearch.neuralsearch.query.NeuralQueryBuilder; + +public class TextSearchIT extends AbstractRestartUpgradeRestTestCase { + +// @Before +// public void setUp() throws Exception { +// super.setUp(); +// updateClusterSettings(); +// prepareModel(); +// } + + @SneakyThrows + public void textSearch() { + waitForClusterHealthGreen(NODES_BWC_CLUSTER); + + + if(isRunningAgainstOldCluster()){ + updateClusterSettings(); + prepareModel(); + initializeBasicIndexIfNotExist(TEST_BASIC_INDEX_NAME); +// String modelId = getDeployedModelId(); +// +// NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( +// TEST_KNN_VECTOR_FIELD_NAME_1, +// TEST_QUERY_TEXT8, +// "", +// modelId, +// 1, +// null, +// null +// ); + + }else { + String modelId = getDeployedModelId(); + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_QUERY_TEXT8, + "", + modelId, + 1, + null, + null + ); + Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); + + assertEquals("1", firstInnerHit.get("_id")); + float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); + assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); + } + + } + + + +} \ No newline at end of file diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java new file mode 100644 index 000000000..089ade9f3 --- /dev/null +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java @@ -0,0 +1,90 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.neuralsearch.bwc; + +import java.util.Locale; +import java.util.Optional; +import org.junit.Before; +import org.opensearch.common.settings.Settings; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; +import org.opensearch.test.rest.OpenSearchRestTestCase; +import static org.opensearch.neuralsearch.TestUtils.CLIENT_TIMEOUT_VALUE; +import static org.opensearch.neuralsearch.TestUtils.OLD_CLUSTER; +import static org.opensearch.neuralsearch.TestUtils.MIXED_CLUSTER; +import static org.opensearch.neuralsearch.TestUtils.UPGRADED_CLUSTER; +import static org.opensearch.neuralsearch.TestUtils.BWC_VERSION; +import static org.opensearch.neuralsearch.TestUtils.ROLLING_UPGRADE_FIRST_ROUND; +import static org.opensearch.neuralsearch.TestUtils.BWCSUITE_CLUSTER; +import static org.opensearch.neuralsearch.TestUtils.NEURAL_SEARCH_BWC_PREFIX; + +public abstract class AbstractRollingUpgradeTestCase extends BaseNeuralSearchIT { + protected String testIndex; + + @Before + protected void setIndex() { + // Creating index name by concatenating "knn-bwc-" prefix with test method name + // for all the tests in this sub-project + testIndex = NEURAL_SEARCH_BWC_PREFIX + getTestName().toLowerCase(Locale.ROOT); + } + + @Override + protected final boolean preserveIndicesUponCompletion() { + return true; + } + + @Override + protected final boolean preserveReposUponCompletion() { + return true; + } + + @Override + protected boolean preserveTemplatesUponCompletion() { + return true; + } + + @Override + protected final Settings restClientSettings() { + return Settings.builder() + .put(super.restClientSettings()) + // increase the timeout here to 90 seconds to handle long waits for a green + // cluster health. the waits for green need to be longer than a minute to + // account for delayed shards + .put(OpenSearchRestTestCase.CLIENT_SOCKET_TIMEOUT, CLIENT_TIMEOUT_VALUE) + .build(); + } + + protected enum ClusterType { + OLD, + MIXED, + UPGRADED; + + public static ClusterType instance(String value) { + switch (value) { + case OLD_CLUSTER: + return OLD; + case MIXED_CLUSTER: + return MIXED; + case UPGRADED_CLUSTER: + return UPGRADED; + default: + throw new IllegalArgumentException("unknown cluster type: " + value); + } + } + } + + protected final ClusterType getClusterType() { + return ClusterType.instance(System.getProperty(BWCSUITE_CLUSTER)); + } + + protected final boolean isFirstMixedRound() { + return Boolean.parseBoolean(System.getProperty(ROLLING_UPGRADE_FIRST_ROUND, "false")); + } + + protected final Optional getBWCVersion() { + return Optional.ofNullable(System.getProperty(BWC_VERSION, null)); + } + +} \ No newline at end of file diff --git a/repositories.gradle b/repositories.gradle new file mode 100644 index 000000000..0ce55eed3 --- /dev/null +++ b/repositories.gradle @@ -0,0 +1,11 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +repositories { + mavenLocal() + maven { url "https://aws.oss.sonatype.org/content/repositories/snapshots" } + mavenCentral() + maven { url "https://plugins.gradle.org/m2/" } +} \ No newline at end of file diff --git a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java index 69791681e..32901cf12 100644 --- a/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java +++ b/src/test/java/org/opensearch/neuralsearch/plugin/NeuralSearchTests.java @@ -11,6 +11,7 @@ import java.util.Map; import java.util.Optional; +import org.opensearch.indices.IndicesService; import org.opensearch.ingest.IngestService; import org.opensearch.ingest.Processor; import org.opensearch.neuralsearch.processor.NeuralQueryEnricherProcessor; @@ -66,7 +67,8 @@ public void testProcessors() { null, mock(IngestService.class), null, - null + null, + mock(IndicesService.class) ); Map processors = plugin.getProcessors(processorParams); assertNotNull(processors); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java index 7e7660457..924eded0a 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java @@ -5,9 +5,11 @@ package org.opensearch.neuralsearch.processor; -import static org.opensearch.neuralsearch.TestUtils.createRandomVector; +import static org.opensearch.neuralsearch.TestUtils.TEST_BASIC_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.ingest_pipeline; +import static org.opensearch.neuralsearch.TestUtils.search_pipeline; -import java.util.Collections; import java.util.Map; import lombok.SneakyThrows; @@ -15,22 +17,10 @@ import org.junit.After; import org.junit.Before; import org.opensearch.common.settings.Settings; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; -import com.google.common.primitives.Floats; - public class NeuralQueryEnricherProcessorIT extends BaseNeuralSearchIT { - - private static final String index = "my-nlp-index"; - private static final String search_pipeline = "search-pipeline"; - private static final String ingest_pipeline = "nlp-pipeline"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; - private static final int TEST_DIMENSION = 768; - private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; - private final float[] testVector = createRandomVector(TEST_DIMENSION); - @Before public void setUp() throws Exception { super.setUp(); @@ -44,40 +34,23 @@ public void tearDown() { super.tearDown(); deleteSearchPipeline(search_pipeline); findDeployedModels().forEach(this::deleteModel); - deleteIndex(index); + deleteIndex(TEST_BASIC_INDEX_NAME); } @SneakyThrows public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() { - initializeIndexIfNotExist(); + initializeBasicIndexIfNotExist(TEST_BASIC_INDEX_NAME); String modelId = getDeployedModelId(); createSearchRequestProcessor(modelId, search_pipeline); createPipelineProcessor(modelId, ingest_pipeline); - updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + updateIndexSettings(TEST_BASIC_INDEX_NAME, Settings.builder().put("index.search.default_pipeline", search_pipeline)); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); neuralQueryBuilder.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1); neuralQueryBuilder.queryText("Hello World"); neuralQueryBuilder.k(1); - Map response = search(index, neuralQueryBuilder, 2); + Map response = search(TEST_BASIC_INDEX_NAME, neuralQueryBuilder, 2); assertFalse(response.isEmpty()); } - - @SneakyThrows - private void initializeIndexIfNotExist() { - if (index.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(index)) { - prepareKnnIndex( - index, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - index, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()) - ); - assertEquals(1, getDocCount(index)); - } - } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java index 86e75f736..b3f5ee7e6 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java @@ -5,11 +5,19 @@ package org.opensearch.neuralsearch.processor; -import static org.opensearch.neuralsearch.TestUtils.createRandomVector; +import static org.opensearch.neuralsearch.TestUtils.RELATION_EQUAL_TO; +import static org.opensearch.neuralsearch.TestUtils.SEARCH_PIPELINE; +import static org.opensearch.neuralsearch.TestUtils.TEST_DOC_TEXT1; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT3; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT4; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT6; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT7; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_1; -import java.io.IOException; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; @@ -23,38 +31,11 @@ import org.junit.Before; import org.opensearch.index.query.QueryBuilders; import org.opensearch.index.query.TermQueryBuilder; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; -import com.google.common.primitives.Floats; - public class NormalizationProcessorIT extends BaseNeuralSearchIT { - private static final String TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME = "test-neural-multi-doc-one-shard-index"; - private static final String TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME = "test-neural-multi-doc-three-shards-index"; - private static final String TEST_QUERY_TEXT3 = "hello"; - private static final String TEST_QUERY_TEXT4 = "place"; - private static final String TEST_QUERY_TEXT6 = "notexistingword"; - private static final String TEST_QUERY_TEXT7 = "notexistingwordtwo"; - private static final String TEST_DOC_TEXT1 = "Hello world"; - private static final String TEST_DOC_TEXT2 = "Hi to this place"; - private static final String TEST_DOC_TEXT3 = "We would like to welcome everyone"; - private static final String TEST_DOC_TEXT4 = "Hello, I'm glad to you see you pal"; - private static final String TEST_DOC_TEXT5 = "Say hello and enter my friend"; - private static final String TEST_DOC_TEXT6 = "This tale grew in the telling"; - private static final String TEST_DOC_TEXT7 = "They do not and did not understand or like machines"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; - private static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field-1"; - private static final String TEST_TEXT_FIELD_NAME_2 = "test-text-field-2"; - private static final int TEST_DIMENSION = 768; - private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; - private static final String SEARCH_PIPELINE = "phase-results-pipeline"; - private final float[] testVector1 = createRandomVector(TEST_DIMENSION); - private final float[] testVector2 = createRandomVector(TEST_DIMENSION); - private final float[] testVector3 = createRandomVector(TEST_DIMENSION); - private final float[] testVector4 = createRandomVector(TEST_DIMENSION); - private final static String RELATION_EQUAL_TO = "eq"; @Before public void setUp() throws Exception { @@ -90,7 +71,7 @@ public void tearDown() { */ @SneakyThrows public void testResultProcessor_whenOneShardAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); String modelId = getDeployedModelId(); @@ -133,7 +114,7 @@ public void testResultProcessor_whenOneShardAndQueryMatches_thenSuccessful() { */ @SneakyThrows public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipelineWithDefaultResultsPostProcessor(SEARCH_PIPELINE); String modelId = getDeployedModelId(); @@ -164,7 +145,7 @@ public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSu @SneakyThrows public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); String modelId = getDeployedModelId(); int totalExpectedDocQty = 6; @@ -224,7 +205,7 @@ public void testResultProcessor_whenMultipleShardsAndQueryMatches_thenSuccessful @SneakyThrows public void testResultProcessor_whenMultipleShardsAndNoMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); @@ -243,7 +224,7 @@ public void testResultProcessor_whenMultipleShardsAndNoMatches_thenSuccessful() @SneakyThrows public void testResultProcessor_whenMultipleShardsAndPartialMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); @@ -261,110 +242,6 @@ public void testResultProcessor_whenMultipleShardsAndPartialMatches_thenSuccessf assertQueryResults(searchResponseAsMap, 4, true, Range.between(0.33f, 1.0f)); } - private void initializeIndexIfNotExist(String indexName) throws IOException { - if (TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), - 1 - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT1) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector2).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector3).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT2) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "4", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT3) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "5", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT4) - ); - assertEquals(5, getDocCount(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)); - } - - if (TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), - 3 - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()), - List.of(TEST_TEXT_FIELD_NAME_1, TEST_TEXT_FIELD_NAME_2), - List.of(TEST_DOC_TEXT1, TEST_DOC_TEXT6) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector2).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector3).toArray()), - List.of(TEST_TEXT_FIELD_NAME_1, TEST_TEXT_FIELD_NAME_2), - List.of(TEST_DOC_TEXT2, TEST_DOC_TEXT7) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "4", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT3) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "5", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT4) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "6", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT5) - ); - assertEquals(6, getDocCount(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME)); - } - } - private List> getNestedHits(Map searchResponseAsMap) { Map hitsMap = (Map) searchResponseAsMap.get("hits"); return (List>) hitsMap.get("hits"); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java index 4993df7fb..8d3327da0 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java @@ -7,13 +7,22 @@ import static org.hamcrest.Matchers.allOf; import static org.hamcrest.Matchers.containsString; +import static org.opensearch.neuralsearch.TestUtils.GEOMETRIC_MEAN_COMBINATION_METHOD; +import static org.opensearch.neuralsearch.TestUtils.HARMONIC_MEAN_COMBINATION_METHOD; +import static org.opensearch.neuralsearch.TestUtils.L2_NORMALIZATION_METHOD; +import static org.opensearch.neuralsearch.TestUtils.SEARCH_PIPELINE; +import static org.opensearch.neuralsearch.TestUtils.TEST_DOC_TEXT1; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT3; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT4; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT7; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_1; import static org.opensearch.neuralsearch.TestUtils.assertHybridSearchResults; import static org.opensearch.neuralsearch.TestUtils.assertWeightedScores; -import static org.opensearch.neuralsearch.TestUtils.createRandomVector; -import java.io.IOException; import java.util.Arrays; -import java.util.Collections; import java.util.Map; import lombok.SneakyThrows; @@ -22,37 +31,11 @@ import org.junit.Before; import org.opensearch.client.ResponseException; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; -import com.google.common.primitives.Floats; - public class ScoreCombinationIT extends BaseNeuralSearchIT { - private static final String TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME = "test-neural-multi-doc-one-shard-index"; - private static final String TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME = "test-neural-multi-doc-three-shards-index"; - private static final String TEST_QUERY_TEXT3 = "hello"; - private static final String TEST_QUERY_TEXT4 = "place"; - private static final String TEST_QUERY_TEXT7 = "notexistingwordtwo"; - private static final String TEST_DOC_TEXT1 = "Hello world"; - private static final String TEST_DOC_TEXT2 = "Hi to this place"; - private static final String TEST_DOC_TEXT3 = "We would like to welcome everyone"; - private static final String TEST_DOC_TEXT4 = "Hello, I'm glad to you see you pal"; - private static final String TEST_DOC_TEXT5 = "Say hello and enter my friend"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; - private static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field-1"; - private static final int TEST_DIMENSION = 768; - private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; - private static final String SEARCH_PIPELINE = "phase-results-pipeline"; - private final float[] testVector1 = createRandomVector(TEST_DIMENSION); - private final float[] testVector2 = createRandomVector(TEST_DIMENSION); - private final float[] testVector3 = createRandomVector(TEST_DIMENSION); - private final float[] testVector4 = createRandomVector(TEST_DIMENSION); - - private static final String L2_NORMALIZATION_METHOD = "l2"; - private static final String HARMONIC_MEAN_COMBINATION_METHOD = "harmonic_mean"; - private static final String GEOMETRIC_MEAN_COMBINATION_METHOD = "geometric_mean"; @Before public void setUp() throws Exception { @@ -93,7 +76,7 @@ public void tearDown() { */ @SneakyThrows public void testArithmeticWeightedMean_whenWeightsPassed_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME); // check case when number of weights and sub-queries are same createSearchPipeline( SEARCH_PIPELINE, @@ -203,7 +186,7 @@ public void testArithmeticWeightedMean_whenWeightsPassed_thenSuccessful() { */ @SneakyThrows public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, @@ -269,7 +252,7 @@ public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessf */ @SneakyThrows public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, @@ -314,108 +297,4 @@ public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccess ); assertHybridSearchResults(searchResponseAsMapL2Norm, 5, new float[] { 0.5f, 1.0f }); } - - private void initializeIndexIfNotExist(String indexName) throws IOException { - if (TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), - 1 - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT1) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector2).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector3).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT2) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "4", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT3) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "5", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT4) - ); - assertEquals(5, getDocCount(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)); - } - - if (TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), - 3 - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT1) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector2).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector3).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT2) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "4", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT3) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "5", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT4) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, - "6", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT5) - ); - assertEquals(6, getDocCount(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME)); - } - } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java index aa133c44d..192d0b37b 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java @@ -5,12 +5,18 @@ package org.opensearch.neuralsearch.processor; +import static org.opensearch.neuralsearch.TestUtils.GEOMETRIC_MEAN_COMBINATION_METHOD; +import static org.opensearch.neuralsearch.TestUtils.HARMONIC_MEAN_COMBINATION_METHOD; +import static org.opensearch.neuralsearch.TestUtils.L2_NORMALIZATION_METHOD; +import static org.opensearch.neuralsearch.TestUtils.SEARCH_PIPELINE; +import static org.opensearch.neuralsearch.TestUtils.TEST_DOC_TEXT1; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT3; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_1; import static org.opensearch.neuralsearch.TestUtils.assertHybridSearchResults; -import static org.opensearch.neuralsearch.TestUtils.createRandomVector; -import java.io.IOException; import java.util.Arrays; -import java.util.Collections; import java.util.Map; import lombok.SneakyThrows; @@ -18,33 +24,11 @@ import org.junit.After; import org.junit.Before; import org.opensearch.index.query.QueryBuilders; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; -import com.google.common.primitives.Floats; - public class ScoreNormalizationIT extends BaseNeuralSearchIT { - private static final String TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME = "test-neural-multi-doc-one-shard-index"; - private static final String TEST_QUERY_TEXT3 = "hello"; - private static final String TEST_DOC_TEXT1 = "Hello world"; - private static final String TEST_DOC_TEXT2 = "Hi to this place"; - private static final String TEST_DOC_TEXT3 = "We would like to welcome everyone"; - private static final String TEST_DOC_TEXT4 = "Hello, I'm glad to you see you pal"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; - private static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field-1"; - private static final int TEST_DIMENSION = 768; - private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; - private static final String SEARCH_PIPELINE = "phase-results-pipeline"; - private final float[] testVector1 = createRandomVector(TEST_DIMENSION); - private final float[] testVector2 = createRandomVector(TEST_DIMENSION); - private final float[] testVector3 = createRandomVector(TEST_DIMENSION); - private final float[] testVector4 = createRandomVector(TEST_DIMENSION); - - private static final String L2_NORMALIZATION_METHOD = "l2"; - private static final String HARMONIC_MEAN_COMBINATION_METHOD = "harmonic_mean"; - private static final String GEOMETRIC_MEAN_COMBINATION_METHOD = "geometric_mean"; @Before public void setUp() throws Exception { @@ -86,7 +70,7 @@ public boolean isUpdateClusterSettings() { */ @SneakyThrows public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, L2_NORMALIZATION_METHOD, @@ -112,7 +96,7 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, L2_NORMALIZATION_METHOD, @@ -137,7 +121,7 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, L2_NORMALIZATION_METHOD, @@ -181,7 +165,7 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { */ @SneakyThrows public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, @@ -207,7 +191,7 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, @@ -232,7 +216,7 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { deleteSearchPipeline(SEARCH_PIPELINE); - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME); createSearchPipeline( SEARCH_PIPELINE, DEFAULT_NORMALIZATION_METHOD, @@ -255,53 +239,4 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { ); assertHybridSearchResults(searchResponseAsMapGeometricMean, 5, new float[] { 0.6f, 1.0f }); } - - private void initializeIndexIfNotExist(String indexName) throws IOException { - if (TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), - 1 - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT1) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector2).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector3).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT2) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "4", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT3) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, - "5", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector4).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT4) - ); - assertEquals(5, getDocCount(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)); - } - } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java b/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java index 893b47ede..cfa926311 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/SparseEncodingProcessIT.java @@ -18,7 +18,7 @@ import org.opensearch.client.Response; import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.common.xcontent.XContentType; -import org.opensearch.neuralsearch.common.BaseSparseEncodingIT; +import org.opensearch.neuralsearch.BaseSparseEncodingIT; import com.google.common.collect.ImmutableList; diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java index de7a70add..e464ae49b 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java @@ -5,6 +5,8 @@ package org.opensearch.neuralsearch.processor; +import static org.opensearch.neuralsearch.TestUtils.TEXT_EMBEDDING_INDEX_NAME; + import java.nio.file.Files; import java.nio.file.Path; import java.util.Map; @@ -18,14 +20,11 @@ import org.opensearch.client.Response; import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.common.xcontent.XContentType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; import com.google.common.collect.ImmutableList; public class TextEmbeddingProcessorIT extends BaseNeuralSearchIT { - - private static final String INDEX_NAME = "text_embedding_index"; - private static final String PIPELINE_NAME = "pipeline-hybrid"; @After @@ -43,9 +42,9 @@ public void testTextEmbeddingProcessor() throws Exception { String modelId = uploadTextEmbeddingModel(); loadModel(modelId); createPipelineProcessor(modelId, PIPELINE_NAME); - createTextEmbeddingIndex(); + createTextEmbeddingIndex(TEXT_EMBEDDING_INDEX_NAME, PIPELINE_NAME); ingestDocument(); - assertEquals(1, getDocCount(INDEX_NAME)); + assertEquals(1, getDocCount(TEXT_EMBEDDING_INDEX_NAME)); } private String uploadTextEmbeddingModel() throws Exception { @@ -53,14 +52,6 @@ private String uploadTextEmbeddingModel() throws Exception { return uploadModel(requestBody); } - private void createTextEmbeddingIndex() throws Exception { - createIndexWithConfiguration( - INDEX_NAME, - Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), - PIPELINE_NAME - ); - } - private void ingestDocument() throws Exception { String ingestDocument = "{\n" + " \"title\": \"This is a good day\",\n" @@ -78,7 +69,7 @@ private void ingestDocument() throws Exception { Response response = makeRequest( client(), "POST", - INDEX_NAME + "/_doc?refresh", + TEXT_EMBEDDING_INDEX_NAME + "/_doc?refresh", null, toHttpEntity(ingestDocument), ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana")) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java index 055344352..8acc8c6c9 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextImageEmbeddingProcessorIT.java @@ -5,6 +5,8 @@ package org.opensearch.neuralsearch.processor; +import static org.opensearch.neuralsearch.TestUtils.TEXT_IMAGE_EMBEDDING_INDEX_NAME; + import java.nio.file.Files; import java.nio.file.Path; import java.util.Map; @@ -18,7 +20,7 @@ import org.opensearch.client.Response; import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.common.xcontent.XContentType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; import com.google.common.collect.ImmutableList; @@ -28,7 +30,6 @@ */ public class TextImageEmbeddingProcessorIT extends BaseNeuralSearchIT { - private static final String INDEX_NAME = "text_image_embedding_index"; private static final String PIPELINE_NAME = "ingest-pipeline"; @After @@ -42,18 +43,18 @@ public void testEmbeddingProcessor_whenIngestingDocumentWithSourceMatchingTextMa String modelId = uploadModel(); loadModel(modelId); createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); - createTextImageEmbeddingIndex(); + createTextEmbeddingIndex(TEXT_IMAGE_EMBEDDING_INDEX_NAME, PIPELINE_NAME); ingestDocumentWithTextMappedToEmbeddingField(); - assertEquals(1, getDocCount(INDEX_NAME)); + assertEquals(1, getDocCount(TEXT_IMAGE_EMBEDDING_INDEX_NAME)); } public void testEmbeddingProcessor_whenIngestingDocumentWithSourceWithoutMatchingInMapping_thenSuccessful() throws Exception { String modelId = uploadModel(); loadModel(modelId); createPipelineProcessor(modelId, PIPELINE_NAME, ProcessorType.TEXT_IMAGE_EMBEDDING); - createTextImageEmbeddingIndex(); + createTextEmbeddingIndex(TEXT_IMAGE_EMBEDDING_INDEX_NAME, PIPELINE_NAME); ingestDocumentWithoutMappedFields(); - assertEquals(1, getDocCount(INDEX_NAME)); + assertEquals(1, getDocCount(TEXT_IMAGE_EMBEDDING_INDEX_NAME)); } private String uploadModel() throws Exception { @@ -61,14 +62,6 @@ private String uploadModel() throws Exception { return uploadModel(requestBody); } - private void createTextImageEmbeddingIndex() throws Exception { - createIndexWithConfiguration( - INDEX_NAME, - Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), - PIPELINE_NAME - ); - } - private void ingestDocumentWithTextMappedToEmbeddingField() throws Exception { String ingestDocumentBody = "{\n" + " \"title\": \"This is a good day\",\n" @@ -95,7 +88,7 @@ private void ingestDocument(final String ingestDocument) throws Exception { Response response = makeRequest( client(), "POST", - INDEX_NAME + "/_doc?refresh", + TEXT_IMAGE_EMBEDDING_INDEX_NAME + "/_doc?refresh", null, toHttpEntity(ingestDocument), ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana")) diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java index eec6955ff..d2c6884d2 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java @@ -6,11 +6,18 @@ package org.opensearch.neuralsearch.query; import static org.opensearch.neuralsearch.TestUtils.DELTA_FOR_SCORE_ASSERTION; -import static org.opensearch.neuralsearch.TestUtils.createRandomVector; +import static org.opensearch.neuralsearch.TestUtils.RELATION_EQUAL_TO; +import static org.opensearch.neuralsearch.TestUtils.SEARCH_PIPELINE; +import static org.opensearch.neuralsearch.TestUtils.TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT2; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT3; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT4; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT5; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_1; -import java.io.IOException; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; @@ -24,34 +31,9 @@ import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.index.query.TermQueryBuilder; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; - -import com.google.common.primitives.Floats; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; public class HybridQueryIT extends BaseNeuralSearchIT { - private static final String TEST_BASIC_INDEX_NAME = "test-neural-basic-index"; - private static final String TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME = "test-neural-vector-doc-field-index"; - private static final String TEST_MULTI_DOC_INDEX_NAME = "test-neural-multi-doc-index"; - private static final String TEST_QUERY_TEXT = "greetings"; - private static final String TEST_QUERY_TEXT2 = "salute"; - private static final String TEST_QUERY_TEXT3 = "hello"; - private static final String TEST_QUERY_TEXT4 = "place"; - private static final String TEST_QUERY_TEXT5 = "welcome"; - private static final String TEST_DOC_TEXT1 = "Hello world"; - private static final String TEST_DOC_TEXT2 = "Hi to this place"; - private static final String TEST_DOC_TEXT3 = "We would like to welcome everyone"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_2 = "test-knn-vector-2"; - private static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field-1"; - - private static final int TEST_DIMENSION = 768; - private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; - private final float[] testVector1 = createRandomVector(TEST_DIMENSION); - private final float[] testVector2 = createRandomVector(TEST_DIMENSION); - private final float[] testVector3 = createRandomVector(TEST_DIMENSION); - private final static String RELATION_EQUAL_TO = "eq"; - private static final String SEARCH_PIPELINE = "phase-results-pipeline"; @Before public void setUp() throws Exception { @@ -117,7 +99,7 @@ protected boolean preserveClusterUponCompletion() { */ @SneakyThrows public void testComplexQuery_whenMultipleSubqueries_thenSuccessful() { - initializeIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME); TermQueryBuilder termQueryBuilder1 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); TermQueryBuilder termQueryBuilder2 = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT4); @@ -161,7 +143,7 @@ public void testComplexQuery_whenMultipleSubqueries_thenSuccessful() { @SneakyThrows public void testNoMatchResults_whenOnlyTermSubQueryWithoutMatch_thenEmptyResult() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); @@ -188,86 +170,6 @@ public void testNoMatchResults_whenOnlyTermSubQueryWithoutMatch_thenEmptyResult( assertEquals(RELATION_EQUAL_TO, total.get("relation")); } - private void initializeIndexIfNotExist(String indexName) throws IOException { - if (TEST_BASIC_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_INDEX_NAME)) { - prepareKnnIndex( - TEST_BASIC_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_BASIC_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()) - ); - assertEquals(1, getDocCount(TEST_BASIC_INDEX_NAME)); - } - if (TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME)) { - prepareKnnIndex( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - List.of( - new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE), - new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_2, TEST_DIMENSION, TEST_SPACE_TYPE) - ) - ); - addKnnDoc( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - "1", - List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), - List.of(Floats.asList(testVector1).toArray(), Floats.asList(testVector1).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT1) - ); - addKnnDoc( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - "2", - List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), - List.of(Floats.asList(testVector2).toArray(), Floats.asList(testVector2).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT2) - ); - addKnnDoc( - TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, - "3", - List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), - List.of(Floats.asList(testVector3).toArray(), Floats.asList(testVector3).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT3) - ); - assertEquals(3, getDocCount(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME)); - } - - if (TEST_MULTI_DOC_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector1).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT1) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector2).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector3).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_DOC_TEXT2) - ); - assertEquals(3, getDocCount(TEST_MULTI_DOC_INDEX_NAME)); - } - } - private List> getNestedHits(Map searchResponseAsMap) { Map hitsMap = (Map) searchResponseAsMap.get("hits"); return (List>) hitsMap.get("hits"); diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java index 6e089500b..0ac90eddc 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java @@ -5,11 +5,22 @@ package org.opensearch.neuralsearch.query; -import static org.opensearch.neuralsearch.TestUtils.createRandomVector; +import static org.opensearch.neuralsearch.TestUtils.TEST_BASIC_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_IMAGE_TEXT; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_2; +import static org.opensearch.neuralsearch.TestUtils.TEST_KNN_VECTOR_FIELD_NAME_NESTED; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_DOC_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_VECTOR_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_NESTED_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT8; +import static org.opensearch.neuralsearch.TestUtils.TEST_SPACE_TYPE; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_3; import static org.opensearch.neuralsearch.TestUtils.objectToFloat; +import static org.opensearch.neuralsearch.TestUtils.testVector; -import java.util.Collections; -import java.util.List; import java.util.Map; import lombok.SneakyThrows; @@ -19,27 +30,9 @@ import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; -import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.common.BaseNeuralSearchIT; - -import com.google.common.primitives.Floats; +import org.opensearch.neuralsearch.BaseNeuralSearchIT; public class NeuralQueryIT extends BaseNeuralSearchIT { - private static final String TEST_BASIC_INDEX_NAME = "test-neural-basic-index"; - private static final String TEST_MULTI_VECTOR_FIELD_INDEX_NAME = "test-neural-multi-vector-field-index"; - private static final String TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME = "test-neural-text-and-vector-field-index"; - private static final String TEST_NESTED_INDEX_NAME = "test-neural-nested-index"; - private static final String TEST_MULTI_DOC_INDEX_NAME = "test-neural-multi-doc-index"; - private static final String TEST_QUERY_TEXT = "Hello world"; - private static final String TEST_IMAGE_TEXT = "/9j/4AAQSkZJRgABAQAASABIAAD"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_2 = "test-knn-vector-2"; - private static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field"; - private static final String TEST_KNN_VECTOR_FIELD_NAME_NESTED = "nested.knn.field"; - - private static final int TEST_DIMENSION = 768; - private static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; - private final float[] testVector = createRandomVector(TEST_DIMENSION); @Before public void setUp() throws Exception { @@ -75,11 +68,11 @@ public void tearDown() { */ @SneakyThrows public void testBasicQuery() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_BASIC_INDEX_NAME); String modelId = getDeployedModelId(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -90,7 +83,7 @@ public void testBasicQuery() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); } @@ -111,11 +104,11 @@ public void testBasicQuery() { */ @SneakyThrows public void testBoostQuery() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_BASIC_INDEX_NAME); String modelId = getDeployedModelId(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -129,7 +122,7 @@ public void testBoostQuery() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); } @@ -155,12 +148,12 @@ public void testBoostQuery() { */ @SneakyThrows public void testRescoreQuery() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_BASIC_INDEX_NAME); String modelId = getDeployedModelId(); MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); NeuralQueryBuilder rescoreNeuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -172,7 +165,7 @@ public void testRescoreQuery() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); } @@ -203,13 +196,13 @@ public void testRescoreQuery() { */ @SneakyThrows public void testBooleanQuery_withMultipleNeuralQueries() { - initializeIndexIfNotExist(TEST_MULTI_VECTOR_FIELD_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_MULTI_VECTOR_FIELD_INDEX_NAME); String modelId = getDeployedModelId(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralQueryBuilder neuralQueryBuilder1 = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -218,7 +211,7 @@ public void testBooleanQuery_withMultipleNeuralQueries() { ); NeuralQueryBuilder neuralQueryBuilder2 = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_2, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -232,7 +225,7 @@ public void testBooleanQuery_withMultipleNeuralQueries() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); } @@ -261,13 +254,13 @@ public void testBooleanQuery_withMultipleNeuralQueries() { */ @SneakyThrows public void testBooleanQuery_withNeuralAndBM25Queries() { - initializeIndexIfNotExist(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME); String modelId = getDeployedModelId(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -275,7 +268,7 @@ public void testBooleanQuery_withNeuralAndBM25Queries() { null ); - MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_3, TEST_QUERY_TEXT8); boolQueryBuilder.should(neuralQueryBuilder).should(matchQueryBuilder); @@ -283,7 +276,7 @@ public void testBooleanQuery_withNeuralAndBM25Queries() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float minExpectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + float minExpectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); } @@ -307,12 +300,12 @@ public void testBooleanQuery_withNeuralAndBM25Queries() { */ @SneakyThrows public void testNestedQuery() { - initializeIndexIfNotExist(TEST_NESTED_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_NESTED_INDEX_NAME); String modelId = getDeployedModelId(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_NESTED, - TEST_QUERY_TEXT, + TEST_QUERY_TEXT8, "", modelId, 1, @@ -324,7 +317,7 @@ public void testNestedQuery() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); + float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT8); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); } @@ -351,7 +344,7 @@ public void testNestedQuery() { */ @SneakyThrows public void testFilterQuery() { - initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + initializeMultiDocIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); String modelId = getDeployedModelId(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, @@ -387,7 +380,7 @@ public void testFilterQuery() { */ @SneakyThrows public void testMultimodalQuery() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeBasicIndexIfNotExist(TEST_BASIC_INDEX_NAME); String modelId = getDeployedModelId(); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder( TEST_KNN_VECTOR_FIELD_NAME_1, @@ -405,94 +398,4 @@ public void testMultimodalQuery() { float expectedScore = computeExpectedScore(modelId, testVector, TEST_SPACE_TYPE, TEST_QUERY_TEXT); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), 0.0); } - - @SneakyThrows - private void initializeIndexIfNotExist(String indexName) { - if (TEST_BASIC_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_INDEX_NAME)) { - prepareKnnIndex( - TEST_BASIC_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_BASIC_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()) - ); - assertEquals(1, getDocCount(TEST_BASIC_INDEX_NAME)); - } - - if (TEST_MULTI_VECTOR_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_VECTOR_FIELD_INDEX_NAME)) { - prepareKnnIndex( - TEST_MULTI_VECTOR_FIELD_INDEX_NAME, - List.of( - new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE), - new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_2, TEST_DIMENSION, TEST_SPACE_TYPE) - ) - ); - addKnnDoc( - TEST_MULTI_VECTOR_FIELD_INDEX_NAME, - "1", - List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), - List.of(Floats.asList(testVector).toArray(), Floats.asList(testVector).toArray()) - ); - assertEquals(1, getDocCount(TEST_MULTI_VECTOR_FIELD_INDEX_NAME)); - } - - if (TEST_NESTED_INDEX_NAME.equals(indexName) && !indexExists(TEST_NESTED_INDEX_NAME)) { - prepareKnnIndex( - TEST_NESTED_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_NESTED, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_NESTED_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_NESTED), - Collections.singletonList(Floats.asList(testVector).toArray()) - ); - assertEquals(1, getDocCount(TEST_NESTED_INDEX_NAME)); - } - - if (TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)) { - prepareKnnIndex( - TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()), - Collections.singletonList(TEST_TEXT_FIELD_NAME_1), - Collections.singletonList(TEST_QUERY_TEXT) - ); - assertEquals(1, getDocCount(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)); - } - - if (TEST_MULTI_DOC_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_NAME)) { - prepareKnnIndex( - TEST_MULTI_DOC_INDEX_NAME, - Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_NAME, - "1", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_NAME, - "2", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()) - ); - addKnnDoc( - TEST_MULTI_DOC_INDEX_NAME, - "3", - Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), - Collections.singletonList(Floats.asList(testVector).toArray()) - ); - assertEquals(3, getDocCount(TEST_MULTI_DOC_INDEX_NAME)); - } - } } diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java index 0a2463ecd..71357255c 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralSparseQueryIT.java @@ -5,9 +5,17 @@ package org.opensearch.neuralsearch.query; +import static org.opensearch.neuralsearch.TestUtils.DELTA; +import static org.opensearch.neuralsearch.TestUtils.TEST_BASIC_SPARSE_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_NEURAL_SPARSE_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.TEST_NEURAL_SPARSE_FIELD_NAME_2; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT9; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_3; import static org.opensearch.neuralsearch.TestUtils.objectToFloat; +import static org.opensearch.neuralsearch.TestUtils.testRankFeaturesDoc; -import java.util.List; import java.util.Map; import lombok.SneakyThrows; @@ -18,24 +26,9 @@ import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; -import org.opensearch.neuralsearch.TestUtils; -import org.opensearch.neuralsearch.common.BaseSparseEncodingIT; +import org.opensearch.neuralsearch.BaseSparseEncodingIT; public class NeuralSparseQueryIT extends BaseSparseEncodingIT { - private static final String TEST_BASIC_INDEX_NAME = "test-sparse-basic-index"; - private static final String TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME = "test-sparse-multi-field-index"; - private static final String TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME = "test-sparse-text-and-field-index"; - private static final String TEST_NESTED_INDEX_NAME = "test-sparse-nested-index"; - private static final String TEST_QUERY_TEXT = "Hello world a b"; - private static final String TEST_NEURAL_SPARSE_FIELD_NAME_1 = "test-sparse-encoding-1"; - private static final String TEST_NEURAL_SPARSE_FIELD_NAME_2 = "test-sparse-encoding-2"; - 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 Float DELTA = 1e-5f; - private final Map testRankFeaturesDoc = TestUtils.createRandomTokenWeightMap(TEST_TOKENS); @Before public void setUp() throws Exception { @@ -66,16 +59,16 @@ public void tearDown() { */ @SneakyThrows public void testBasicQueryUsingQueryText() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeSparseIndexIfNotExist(TEST_BASIC_SPARSE_INDEX_NAME); String modelId = getDeployedModelId(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map searchResponseAsMap = search(TEST_BASIC_SPARSE_INDEX_NAME, sparseEncodingQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT9); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @@ -95,17 +88,17 @@ public void testBasicQueryUsingQueryText() { */ @SneakyThrows public void testBoostQuery() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeSparseIndexIfNotExist(TEST_BASIC_SPARSE_INDEX_NAME); String modelId = getDeployedModelId(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId) .boost(2.0f); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, sparseEncodingQueryBuilder, 1); + Map searchResponseAsMap = search(TEST_BASIC_SPARSE_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(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT9); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @@ -131,17 +124,17 @@ public void testBoostQuery() { */ @SneakyThrows public void testRescoreQuery() { - initializeIndexIfNotExist(TEST_BASIC_INDEX_NAME); + initializeSparseIndexIfNotExist(TEST_BASIC_SPARSE_INDEX_NAME); String modelId = getDeployedModelId(); MatchAllQueryBuilder matchAllQueryBuilder = new MatchAllQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId); - Map searchResponseAsMap = search(TEST_BASIC_INDEX_NAME, matchAllQueryBuilder, sparseEncodingQueryBuilder, 1); + Map searchResponseAsMap = search(TEST_BASIC_SPARSE_INDEX_NAME, matchAllQueryBuilder, sparseEncodingQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT9); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @@ -170,15 +163,15 @@ public void testRescoreQuery() { */ @SneakyThrows public void testBooleanQuery_withMultipleSparseEncodingQueries() { - initializeIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); + initializeSparseIndexIfNotExist(TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME); String modelId = getDeployedModelId(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder1 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId); NeuralSparseQueryBuilder sparseEncodingQueryBuilder2 = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_2) - .queryText(TEST_QUERY_TEXT) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId); boolQueryBuilder.should(sparseEncodingQueryBuilder1).should(sparseEncodingQueryBuilder2); @@ -187,7 +180,7 @@ public void testBooleanQuery_withMultipleSparseEncodingQueries() { Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float expectedScore = 2 * computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT9); assertEquals(expectedScore, objectToFloat(firstInnerHit.get("_score")), DELTA); } @@ -216,72 +209,33 @@ public void testBooleanQuery_withMultipleSparseEncodingQueries() { */ @SneakyThrows public void testBooleanQuery_withSparseEncodingAndBM25Queries() { - initializeIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); + initializeSparseIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); String modelId = getDeployedModelId(); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_NEURAL_SPARSE_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId); - MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(TEST_TEXT_FIELD_NAME_3, TEST_QUERY_TEXT9); boolQueryBuilder.should(sparseEncodingQueryBuilder).should(matchQueryBuilder); Map searchResponseAsMap = search(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, boolQueryBuilder, 1); Map firstInnerHit = getFirstInnerHit(searchResponseAsMap); assertEquals("1", firstInnerHit.get("_id")); - float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT); + float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc, TEST_QUERY_TEXT9); assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); } @SneakyThrows public void testBasicQueryUsingQueryText_whenQueryWrongFieldType_thenFail() { - initializeIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); + initializeSparseIndexIfNotExist(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME); String modelId = getDeployedModelId(); - NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_TEXT_FIELD_NAME_1) - .queryText(TEST_QUERY_TEXT) + NeuralSparseQueryBuilder sparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_TEXT_FIELD_NAME_3) + .queryText(TEST_QUERY_TEXT9) .modelId(modelId); expectThrows(ResponseException.class, () -> search(TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME, sparseEncodingQueryBuilder, 1)); } - - @SneakyThrows - protected void initializeIndexIfNotExist(String indexName) { - if (TEST_BASIC_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { - prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1)); - addSparseEncodingDoc(indexName, "1", List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1), List.of(testRankFeaturesDoc)); - assertEquals(1, getDocCount(indexName)); - } - - if (TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { - prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1, TEST_NEURAL_SPARSE_FIELD_NAME_2)); - addSparseEncodingDoc( - indexName, - "1", - List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1, TEST_NEURAL_SPARSE_FIELD_NAME_2), - List.of(testRankFeaturesDoc, testRankFeaturesDoc) - ); - assertEquals(1, getDocCount(indexName)); - } - - if (TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { - prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1)); - addSparseEncodingDoc( - indexName, - "1", - List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1), - List.of(testRankFeaturesDoc), - List.of(TEST_TEXT_FIELD_NAME_1), - List.of(TEST_QUERY_TEXT) - ); - assertEquals(1, getDocCount(indexName)); - } - - if (TEST_NESTED_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { - prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_NESTED)); - addSparseEncodingDoc(indexName, "1", List.of(TEST_NEURAL_SPARSE_FIELD_NAME_NESTED), List.of(testRankFeaturesDoc)); - assertEquals(1, getDocCount(TEST_NESTED_INDEX_NAME)); - } - } } diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java similarity index 69% rename from src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java rename to src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java index 33cdff9a0..b5327ff45 100644 --- a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java @@ -3,21 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.neuralsearch.common; +package org.opensearch.neuralsearch; +import static org.opensearch.neuralsearch.TestUtils.*; import static org.opensearch.neuralsearch.common.VectorUtil.vectorAsListToArray; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; -import java.util.Collections; -import java.util.HashSet; -import java.util.List; -import java.util.Locale; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; -import java.util.Set; +import java.util.*; import java.util.function.Predicate; import java.util.stream.Collectors; @@ -34,11 +28,7 @@ import org.apache.hc.core5.http.io.entity.StringEntity; import org.apache.hc.core5.http.message.BasicHeader; import org.junit.Before; -import org.opensearch.client.Request; -import org.opensearch.client.RequestOptions; -import org.opensearch.client.Response; -import org.opensearch.client.RestClient; -import org.opensearch.client.WarningsHandler; +import org.opensearch.client.*; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentFactory; @@ -49,7 +39,6 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.knn.index.SpaceType; -import org.opensearch.neuralsearch.OpenSearchSecureRestTestCase; import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; import org.opensearch.test.ClusterServiceUtils; import org.opensearch.threadpool.TestThreadPool; @@ -58,6 +47,7 @@ import com.carrotsearch.randomizedtesting.RandomizedTest; import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; import com.google.common.collect.ImmutableList; +import com.google.common.primitives.Floats; @ThreadLeakScope(ThreadLeakScope.Scope.NONE) public abstract class BaseNeuralSearchIT extends OpenSearchSecureRestTestCase { @@ -772,6 +762,283 @@ private String registerModelGroup() { return modelGroupId; } + // Method that waits till the health of nodes in the cluster goes green + public void waitForClusterHealthGreen(String numOfNodes) throws IOException { + Request waitForGreen = new Request("GET", "/_cluster/health"); + waitForGreen.addParameter("wait_for_nodes", numOfNodes); + waitForGreen.addParameter("wait_for_status", "green"); + client().performRequest(waitForGreen); + } + + @SneakyThrows + protected void initializeBasicIndexIfNotExist(String indexName) { + if (TEST_BASIC_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_INDEX_NAME)) { + prepareKnnIndex( + TEST_BASIC_INDEX_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) + ); + addKnnDoc( + TEST_BASIC_INDEX_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector).toArray()) + ); + assertEquals(1, getDocCount(TEST_BASIC_INDEX_NAME)); + } + + if (TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME)) { + prepareKnnIndex( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + List.of( + new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE), + new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_2, TEST_DIMENSION, TEST_SPACE_TYPE) + ) + ); + addKnnDoc( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + "1", + List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), + List.of(Floats.asList(testVector1).toArray(), Floats.asList(testVector1).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT1) + ); + addKnnDoc( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + "2", + List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), + List.of(Floats.asList(testVector2).toArray(), Floats.asList(testVector2).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT2) + ); + addKnnDoc( + TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME, + "3", + List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), + List.of(Floats.asList(testVector3).toArray(), Floats.asList(testVector3).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT3) + ); + assertEquals(3, getDocCount(TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME)); + } + + if (TEST_MULTI_DOC_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_NAME)) { + prepareKnnIndex( + TEST_MULTI_DOC_INDEX_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector1).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT1) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_NAME, + "2", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector2).toArray()) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_NAME, + "3", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector3).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT2) + ); + assertEquals(3, getDocCount(TEST_MULTI_DOC_INDEX_NAME)); + } + + if (TEST_MULTI_VECTOR_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_VECTOR_FIELD_INDEX_NAME)) { + prepareKnnIndex( + TEST_MULTI_VECTOR_FIELD_INDEX_NAME, + List.of( + new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE), + new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_2, TEST_DIMENSION, TEST_SPACE_TYPE) + ) + ); + addKnnDoc( + TEST_MULTI_VECTOR_FIELD_INDEX_NAME, + "1", + List.of(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_KNN_VECTOR_FIELD_NAME_2), + List.of(Floats.asList(testVector).toArray(), Floats.asList(testVector).toArray()) + ); + assertEquals(1, getDocCount(TEST_MULTI_VECTOR_FIELD_INDEX_NAME)); + } + + if (TEST_NESTED_INDEX_NAME.equals(indexName) && !indexExists(TEST_NESTED_INDEX_NAME)) { + prepareKnnIndex( + TEST_NESTED_INDEX_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_NESTED, TEST_DIMENSION, TEST_SPACE_TYPE)) + ); + addKnnDoc( + TEST_NESTED_INDEX_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_NESTED), + Collections.singletonList(Floats.asList(testVector).toArray()) + ); + assertEquals(1, getDocCount(TEST_NESTED_INDEX_NAME)); + } + + if (TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME.equals(indexName) && !indexExists(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)) { + prepareKnnIndex( + TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) + ); + addKnnDoc( + TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_3), + Collections.singletonList(TEST_QUERY_TEXT) + ); + assertEquals(1, getDocCount(TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME)); + } + } + + @SneakyThrows + protected void initializeMultiDocIndexIfNotExist(String indexName) { + if (TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)) { + prepareKnnIndex( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), + 1 + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector1).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT1) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + "2", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector2).toArray()) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + "3", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector3).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT2) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + "4", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector4).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT3) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME, + "5", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector4).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT4) + ); + assertEquals(5, getDocCount(TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME)); + } + + if (TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME.equalsIgnoreCase(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME)) { + prepareKnnIndex( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)), + 3 + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector1).toArray()), + List.of(TEST_TEXT_FIELD_NAME_1, TEST_TEXT_FIELD_NAME_2), + List.of(TEST_DOC_TEXT1, TEST_DOC_TEXT6) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + "2", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector2).toArray()) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + "3", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector3).toArray()), + List.of(TEST_TEXT_FIELD_NAME_1, TEST_TEXT_FIELD_NAME_2), + List.of(TEST_DOC_TEXT2, TEST_DOC_TEXT7) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + "4", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector4).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT3) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + "5", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector4).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT4) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME, + "6", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector4).toArray()), + Collections.singletonList(TEST_TEXT_FIELD_NAME_1), + Collections.singletonList(TEST_DOC_TEXT5) + ); + assertEquals(6, getDocCount(TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME)); + } + + if (TEST_MULTI_DOC_INDEX_NAME.equals(indexName) && !indexExists(TEST_MULTI_DOC_INDEX_NAME)) { + prepareKnnIndex( + TEST_MULTI_DOC_INDEX_NAME, + Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_NAME, + "1", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector).toArray()) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_NAME, + "2", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector).toArray()) + ); + addKnnDoc( + TEST_MULTI_DOC_INDEX_NAME, + "3", + Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1), + Collections.singletonList(Floats.asList(testVector).toArray()) + ); + assertEquals(3, getDocCount(TEST_MULTI_DOC_INDEX_NAME)); + } + + } + + public void createTextEmbeddingIndex(String indexName, String pipelineName) throws Exception { + createIndexWithConfiguration( + indexName, + Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())), + pipelineName + ); + } + /** * Enumeration for types of pipeline processors, used to lookup resources like create * processor request as those are type specific diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseSparseEncodingIT.java b/src/testFixtures/java/org/opensearch/neuralsearch/BaseSparseEncodingIT.java similarity index 67% rename from src/test/java/org/opensearch/neuralsearch/common/BaseSparseEncodingIT.java rename to src/testFixtures/java/org/opensearch/neuralsearch/BaseSparseEncodingIT.java index d0231cfe6..a98e6dcfd 100644 --- a/src/test/java/org/opensearch/neuralsearch/common/BaseSparseEncodingIT.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/BaseSparseEncodingIT.java @@ -3,7 +3,18 @@ * SPDX-License-Identifier: Apache-2.0 */ -package org.opensearch.neuralsearch.common; +package org.opensearch.neuralsearch; + +import static org.opensearch.neuralsearch.TestUtils.TEST_BASIC_SPARSE_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_NESTED_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_NEURAL_SPARSE_FIELD_NAME_1; +import static org.opensearch.neuralsearch.TestUtils.TEST_NEURAL_SPARSE_FIELD_NAME_2; +import static org.opensearch.neuralsearch.TestUtils.TEST_NEURAL_SPARSE_FIELD_NAME_NESTED; +import static org.opensearch.neuralsearch.TestUtils.TEST_QUERY_TEXT9; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME; +import static org.opensearch.neuralsearch.TestUtils.TEST_TEXT_FIELD_NAME_3; +import static org.opensearch.neuralsearch.TestUtils.testRankFeaturesDoc; import java.nio.file.Files; import java.nio.file.Path; @@ -136,4 +147,43 @@ protected Float getFeatureFieldCompressedNumber(Float originNumber) { freqBits = ((int) ((float) freqBits)) << 15; return Float.intBitsToFloat(freqBits); } + + @SneakyThrows + protected void initializeSparseIndexIfNotExist(String indexName) { + if (TEST_BASIC_SPARSE_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { + prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1)); + addSparseEncodingDoc(indexName, "1", List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1), List.of(testRankFeaturesDoc)); + assertEquals(1, getDocCount(indexName)); + } + + if (TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { + prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1, TEST_NEURAL_SPARSE_FIELD_NAME_2)); + addSparseEncodingDoc( + indexName, + "1", + List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1, TEST_NEURAL_SPARSE_FIELD_NAME_2), + List.of(testRankFeaturesDoc, testRankFeaturesDoc) + ); + assertEquals(1, getDocCount(indexName)); + } + + if (TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { + prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1)); + addSparseEncodingDoc( + indexName, + "1", + List.of(TEST_NEURAL_SPARSE_FIELD_NAME_1), + List.of(testRankFeaturesDoc), + List.of(TEST_TEXT_FIELD_NAME_3), + List.of(TEST_QUERY_TEXT9) + ); + assertEquals(1, getDocCount(indexName)); + } + + if (TEST_NESTED_INDEX_NAME.equals(indexName) && !indexExists(indexName)) { + prepareSparseEncodingIndex(indexName, List.of(TEST_NEURAL_SPARSE_FIELD_NAME_NESTED)); + addSparseEncodingDoc(indexName, "1", List.of(TEST_NEURAL_SPARSE_FIELD_NAME_NESTED), List.of(testRankFeaturesDoc)); + assertEquals(1, getDocCount(TEST_NESTED_INDEX_NAME)); + } + } } diff --git a/src/test/java/org/opensearch/neuralsearch/OpenSearchSecureRestTestCase.java b/src/testFixtures/java/org/opensearch/neuralsearch/OpenSearchSecureRestTestCase.java similarity index 100% rename from src/test/java/org/opensearch/neuralsearch/OpenSearchSecureRestTestCase.java rename to src/testFixtures/java/org/opensearch/neuralsearch/OpenSearchSecureRestTestCase.java diff --git a/src/test/java/org/opensearch/neuralsearch/TestUtils.java b/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java similarity index 70% rename from src/test/java/org/opensearch/neuralsearch/TestUtils.java rename to src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java index 385855a2e..58e95e016 100644 --- a/src/test/java/org/opensearch/neuralsearch/TestUtils.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java @@ -27,6 +27,7 @@ import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.knn.index.SpaceType; import org.opensearch.search.SearchHit; import org.opensearch.search.SearchHits; import org.opensearch.search.fetch.FetchSearchResult; @@ -34,8 +35,75 @@ public class TestUtils { - private static final String RELATION_EQUAL_TO = "eq"; + public static final String RELATION_EQUAL_TO = "eq"; public static final float DELTA_FOR_SCORE_ASSERTION = 0.001f; + public static final String CLIENT_TIMEOUT_VALUE = "90s"; + public static final String RESTART_UPGRADE_OLD_CLUSTER = "tests.is_old_cluster"; + public static final String BWC_VERSION = "tests.plugin_bwc_version"; + public static final String NEURAL_SEARCH_BWC_PREFIX = "neuralsearch-bwc-"; + public static final String OLD_CLUSTER = "old_cluster"; + public static final String MIXED_CLUSTER = "mixed_cluster"; + public static final String UPGRADED_CLUSTER = "upgraded_cluster"; + public static final String ROLLING_UPGRADE_FIRST_ROUND = "tests.rest.first_round"; + public static final String BWCSUITE_CLUSTER = "tests.rest.bwcsuite_cluster"; + public static final String NODES_BWC_CLUSTER = "3"; + public static final String TEST_MULTI_DOC_INDEX_ONE_SHARD_NAME = "test-neural-multi-doc-one-shard-index"; + public static final String TEST_MULTI_DOC_INDEX_THREE_SHARDS_NAME = "test-neural-multi-doc-three-shards-index"; + public static final String TEST_QUERY_TEXT3 = "hello"; + public static final String TEST_QUERY_TEXT4 = "place"; + public static final String TEST_QUERY_TEXT6 = "notexistingword"; + public static final String TEST_QUERY_TEXT7 = "notexistingwordtwo"; + public static final String TEST_DOC_TEXT1 = "Hello world"; + public static final String TEST_DOC_TEXT2 = "Hi to this place"; + public static final String TEST_DOC_TEXT3 = "We would like to welcome everyone"; + public static final String TEST_DOC_TEXT4 = "Hello, I'm glad to you see you pal"; + public static final String TEST_DOC_TEXT5 = "Say hello and enter my friend"; + public static final String TEST_DOC_TEXT6 = "This tale grew in the telling"; + public static final String TEST_DOC_TEXT7 = "They do not and did not understand or like machines"; + public static final String TEST_KNN_VECTOR_FIELD_NAME_1 = "test-knn-vector-1"; + public static final String TEST_TEXT_FIELD_NAME_1 = "test-text-field-1"; + public static final String TEST_TEXT_FIELD_NAME_2 = "test-text-field-2"; + public static final int TEST_DIMENSION = 768; + public static final SpaceType TEST_SPACE_TYPE = SpaceType.L2; + public static final String search_pipeline = "search-pipeline"; + public static final String ingest_pipeline = "nlp-pipeline"; + public static final float[] testVector = createRandomVector(TEST_DIMENSION); + public static final float[] testVector1 = createRandomVector(TEST_DIMENSION); + public static final float[] testVector2 = createRandomVector(TEST_DIMENSION); + public static final float[] testVector3 = createRandomVector(TEST_DIMENSION); + public static final float[] testVector4 = createRandomVector(TEST_DIMENSION); + + public static final String L2_NORMALIZATION_METHOD = "l2"; + public static final String HARMONIC_MEAN_COMBINATION_METHOD = "harmonic_mean"; + public static final String GEOMETRIC_MEAN_COMBINATION_METHOD = "geometric_mean"; + public static final String SEARCH_PIPELINE = "phase-results-pipeline"; + public static final String TEST_QUERY_TEXT = "greetings"; + public static final String TEST_QUERY_TEXT2 = "salute"; + public static final String TEST_QUERY_TEXT5 = "welcome"; + public static final String TEST_KNN_VECTOR_FIELD_NAME_NESTED = "nested.knn.field"; + public static final String TEST_KNN_VECTOR_FIELD_NAME_2 = "test-knn-vector-2"; + public static final String TEST_IMAGE_TEXT = "/9j/4AAQSkZJRgABAQAASABIAAD"; + public static final String TEST_TEXT_FIELD_NAME_3 = "test-text-field"; + public static final String TEST_QUERY_TEXT8 = "Hello world"; + public static final String TEST_QUERY_TEXT9 = "Hello world a b"; + public static final String TEST_NEURAL_SPARSE_FIELD_NAME_1 = "test-sparse-encoding-1"; + public static final String TEST_NEURAL_SPARSE_FIELD_NAME_2 = "test-sparse-encoding-2"; + public static final String TEST_NEURAL_SPARSE_FIELD_NAME_NESTED = "nested.neural_sparse.field"; + public static final List TEST_TOKENS = List.of("hello", "world", "a", "b", "c"); + public static final Float DELTA = 1e-5f; + public static final Map testRankFeaturesDoc = TestUtils.createRandomTokenWeightMap(TEST_TOKENS); + public static final String TEST_BASIC_INDEX_NAME = "test-neural-basic-index"; + public static final String TEST_BASIC_VECTOR_DOC_FIELD_INDEX_NAME = "test-neural-vector-doc-field-index"; + public static final String TEST_MULTI_DOC_INDEX_NAME = "test-neural-multi-doc-index"; + public static final String TEST_MULTI_VECTOR_FIELD_INDEX_NAME = "test-neural-multi-vector-field-index"; + public static final String TEST_TEXT_AND_VECTOR_FIELD_INDEX_NAME = "test-neural-text-and-vector-field-index"; + public static final String TEST_NESTED_INDEX_NAME = "test-neural-nested-index"; + public static final String TEST_BASIC_SPARSE_INDEX_NAME = "test-sparse-basic-index"; + public static final String TEST_MULTI_NEURAL_SPARSE_FIELD_INDEX_NAME = "test-sparse-multi-field-index"; + public static final String TEST_TEXT_AND_NEURAL_SPARSE_FIELD_INDEX_NAME = "test-sparse-text-and-field-index"; + public static final String TEST_NESTED_SPARSE_INDEX_NAME = "test-sparse-nested-index"; + public static final String TEXT_EMBEDDING_INDEX_NAME = "text_embedding_index"; + public static final String TEXT_IMAGE_EMBEDDING_INDEX_NAME = "text_image_embedding_index"; /** * Convert an xContentBuilder to a map