1010import org .elasticsearch .action .DocWriteResponse ;
1111import org .elasticsearch .action .admin .indices .refresh .RefreshRequest ;
1212import org .elasticsearch .action .search .SearchRequest ;
13+ import org .elasticsearch .cluster .metadata .IndexMetadata ;
14+ import org .elasticsearch .common .io .stream .NamedWriteableRegistry ;
1315import org .elasticsearch .common .settings .Settings ;
1416import org .elasticsearch .core .TimeValue ;
1517import org .elasticsearch .index .IndexVersion ;
16- import org .elasticsearch .index .mapper .InferenceMetadataFieldsMapper ;
18+ import org .elasticsearch .index .mapper .vectors .DenseVectorFieldMapper ;
19+ import org .elasticsearch .index .mapper .vectors .DenseVectorFieldMapperTestUtils ;
20+ import org .elasticsearch .index .query .QueryBuilder ;
21+ import org .elasticsearch .inference .SimilarityMeasure ;
1722import org .elasticsearch .license .LicenseSettings ;
1823import org .elasticsearch .plugins .Plugin ;
1924import org .elasticsearch .search .builder .SearchSourceBuilder ;
2227import org .elasticsearch .test .index .IndexVersionUtils ;
2328import org .elasticsearch .xcontent .XContentBuilder ;
2429import org .elasticsearch .xcontent .XContentFactory ;
30+ import org .elasticsearch .xpack .core .ml .inference .MlInferenceNamedXContentProvider ;
31+ import org .elasticsearch .xpack .core .ml .search .SparseVectorQueryBuilder ;
2532import org .elasticsearch .xpack .inference .LocalStateInferencePlugin ;
2633import org .elasticsearch .xpack .inference .Utils ;
2734import org .elasticsearch .xpack .inference .mock .TestSparseInferenceServiceExtension ;
2835import org .elasticsearch .xpack .inference .queries .SemanticQueryBuilder ;
2936import org .elasticsearch .xpack .inference .registry .ModelRegistry ;
3037import org .junit .Before ;
3138
32- import java .io . IOException ;
39+ import java .util . ArrayList ;
3340import java .util .Collection ;
34- import java .util .HashMap ;
3541import java .util .List ;
3642import java .util .Locale ;
3743import java .util .Map ;
3844import java .util .Set ;
45+ import java .util .function .Function ;
3946import java .util .stream .Collectors ;
4047
4148import static org .elasticsearch .test .hamcrest .ElasticsearchAssertions .assertAcked ;
4754public class SemanticTextIndexVersionIT extends ESIntegTestCase {
4855 private static final IndexVersion SEMANTIC_TEXT_INTRODUCED_VERSION = IndexVersion .fromId (8512000 );
4956 private static final double PERCENTAGE_TO_TEST = 0.5 ;
50- private Set <IndexVersion > availableVersions ;
57+ private static final int MAXIMUM_NUMBER_OF_VERSIONS_TO_TEST = 25 ;
58+ private List <IndexVersion > selectedVersions ;
5159
5260 @ Before
5361 public void setup () throws Exception {
5462 ModelRegistry modelRegistry = internalCluster ().getCurrentMasterNodeInstance (ModelRegistry .class );
63+ DenseVectorFieldMapper .ElementType elementType = randomFrom (DenseVectorFieldMapper .ElementType .values ());
64+ // dot product means that we need normalized vectors; it's not worth doing that in this test
65+ SimilarityMeasure similarity = randomValueOtherThan (
66+ SimilarityMeasure .DOT_PRODUCT ,
67+ () -> randomFrom (DenseVectorFieldMapperTestUtils .getSupportedSimilarities (elementType ))
68+ );
69+ int dimensions = DenseVectorFieldMapperTestUtils .randomCompatibleDimensions (elementType , 100 );
5570 Utils .storeSparseModel (modelRegistry );
56- availableVersions = IndexVersionUtils .allReleasedVersions ()
71+ Utils .storeDenseModel (modelRegistry , dimensions , similarity , elementType );
72+
73+ Set <IndexVersion > availableVersions = IndexVersionUtils .allReleasedVersions ()
5774 .stream ()
5875 .filter (indexVersion -> indexVersion .after (SEMANTIC_TEXT_INTRODUCED_VERSION ))
5976 .collect (Collectors .toSet ());
77+
78+ Function <Set <IndexVersion >, Integer > determineNumberOfVersionsToTest = versions -> {
79+ int totalVersions = versions .size ();
80+ int percentageTestSize = (int ) Math .ceil (totalVersions * PERCENTAGE_TO_TEST );
81+
82+ return totalVersions < MAXIMUM_NUMBER_OF_VERSIONS_TO_TEST
83+ ? totalVersions
84+ : Math .min (percentageTestSize , MAXIMUM_NUMBER_OF_VERSIONS_TO_TEST );
85+ };
86+
87+ int versionsCount = determineNumberOfVersionsToTest .apply (availableVersions );
88+ selectedVersions = randomSubsetOf (versionsCount , availableVersions );
6089 }
6190
6291 @ Override
@@ -71,44 +100,24 @@ protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
71100
72101 @ Override
73102 protected Collection <Class <? extends Plugin >> nodePlugins () {
74- return List .of (LocalStateInferencePlugin .class );
103+ return List .of (LocalStateInferencePlugin .class , FakeMlPlugin . class , FakeSemanticQueryBuilderPlugin . class );
75104 }
76105
77106 /**
78107 * Generate settings for an index with a specific version.
79108 */
80109 private Settings getIndexSettingsWithVersion (IndexVersion version ) {
81- return Settings .builder ().put (indexSettings ()).put ("index.version.created" , version ).build ();
82- }
83-
84- /**
85- * Creates a subset of indices with different versions for testing.
86- *
87- * @return Map of created indices with their versions
88- */
89- protected Map <String , IndexVersion > createRandomVersionIndices () throws IOException {
90- int versionsCount = (int ) Math .ceil (availableVersions .size () * PERCENTAGE_TO_TEST );
91- List <IndexVersion > selectedVersions = randomSubsetOf (versionsCount , availableVersions );
92- Map <String , IndexVersion > result = new HashMap <>();
93-
94- for (int i = 0 ; i < selectedVersions .size (); i ++) {
95- String indexName = "test_semantic" + "_" + i ;
96- IndexVersion version = selectedVersions .get (i );
97- createIndex (indexName , getIndexSettingsWithVersion (version ));
98- result .put (indexName , version );
99- }
100-
101- return result ;
110+ return Settings .builder ().put (indexSettings ()).put (IndexMetadata .SETTING_VERSION_CREATED , version ).build ();
102111 }
103112
104113 /**
105114 * This test creates an index, ingests data, and performs searches (including highlighting when applicable)
106115 * for a selected subset of index versions.
107116 */
108117 public void testSemanticText () throws Exception {
109- Map < String , IndexVersion > indices = createRandomVersionIndices ();
110- for ( String indexName : indices . keySet ()) {
111- IndexVersion version = indices . get (indexName );
118+ for ( IndexVersion version : selectedVersions ) {
119+ String indexName = "test_semantic_" + randomAlphaOfLength ( 5 ). toLowerCase ( Locale . ROOT );
120+ createIndex (indexName , getIndexSettingsWithVersion ( version ) );
112121
113122 // Test index creation
114123 assertTrue ("Index " + indexName + " should exist" , indexExists (indexName ));
@@ -164,17 +173,36 @@ public void testSemanticText() throws Exception {
164173 .getIndexToSettings ()
165174 .get (indexName );
166175
167- // Semantic Search with highlighter only available from 8.18 and 9.0
168- if (InferenceMetadataFieldsMapper .isEnabled (settings )) {
169- SearchSourceBuilder sourceHighlighterBuilder = new SearchSourceBuilder ().query (
170- new SemanticQueryBuilder ("semantic_field" , "inference" )
171- ).highlighter (new HighlightBuilder ().field ("semantic_field" )).trackTotalHits (true );
172-
173- assertResponse (client ().search (new SearchRequest (indexName ).source (sourceHighlighterBuilder )), response -> {
174- assertHighlight (response , 0 , "semantic_field" , 0 , 2 , equalTo ("inference test" ));
175- assertHighlight (response , 0 , "semantic_field" , 1 , 2 , equalTo ("another inference test" ));
176- });
177- }
176+ // Semantic Search with highlighter
177+ SearchSourceBuilder sourceHighlighterBuilder = new SearchSourceBuilder ().query (
178+ new SemanticQueryBuilder ("semantic_field" , "inference" )
179+ ).highlighter (new HighlightBuilder ().field ("semantic_field" )).trackTotalHits (true );
180+
181+ assertResponse (client ().search (new SearchRequest (indexName ).source (sourceHighlighterBuilder )), response -> {
182+ assertHighlight (response , 0 , "semantic_field" , 0 , 2 , equalTo ("inference test" ));
183+ assertHighlight (response , 0 , "semantic_field" , 1 , 2 , equalTo ("another inference test" ));
184+ });
185+
186+ beforeIndexDeletion ();
187+ assertAcked (client ().admin ().indices ().prepareDelete (indexName ));
188+ }
189+ }
190+
191+ public static class FakeMlPlugin extends Plugin {
192+ @ Override
193+ public List <NamedWriteableRegistry .Entry > getNamedWriteables () {
194+ return new MlInferenceNamedXContentProvider ().getNamedWriteables ();
195+ }
196+ }
197+
198+ public static class FakeSemanticQueryBuilderPlugin extends Plugin {
199+ @ Override
200+ public List <NamedWriteableRegistry .Entry > getNamedWriteables () {
201+ List <NamedWriteableRegistry .Entry > namedWriteables = new ArrayList <>();
202+ namedWriteables .add (
203+ new NamedWriteableRegistry .Entry (QueryBuilder .class , SparseVectorQueryBuilder .NAME , SparseVectorQueryBuilder ::new )
204+ );
205+ return namedWriteables ;
178206 }
179207 }
180208}
0 commit comments