Merge branch 'main' of github.com:elastic/elasticsearch into cs-request-params

Author: jonathan-buttner

docs/changelog/126866.yaml

@@ -0,0 +1,5 @@
+pr: 126866
+summary: Add recursive chunker
+area: Machine Learning
+type: enhancement
+issues: []

libs/simdvec/src/test/java/org/elasticsearch/simdvec/internal/vectorization/ES91OSQVectorScorerTests.java

@@ -125,12 +125,21 @@ public void testScore() throws Exception {
                             scores2
                         );
                         for (int j = 0; j < ES91OSQVectorsScorer.BULK_SIZE; j++) {
-                            if (scores1[j] > (maxDims * Short.MAX_VALUE)) {
-                                int diff = (int) (scores1[j] - scores2[j]);
-                                assertThat("defaultScores: " + scores1[j] + " bulkScores: " + scores2[j], Math.abs(diff), lessThan(65));
-                            } else if (scores1[j] > (maxDims * Byte.MAX_VALUE)) {
-                                int diff = (int) (scores1[j] - scores2[j]);
-                                assertThat("defaultScores: " + scores1[j] + " bulkScores: " + scores2[j], Math.abs(diff), lessThan(9));
+                            if (scores1[j] == scores2[j]) {
+                                continue;
+                            }
+                            if (scores1[j] > (maxDims * Byte.MAX_VALUE)) {
+                                float diff = Math.abs(scores1[j] - scores2[j]);
+                                assertThat(
+                                    "defaultScores: " + scores1[j] + " bulkScores: " + scores2[j],
+                                    diff / scores1[j],
+                                    lessThan(1e-5f)
+                                );
+                                assertThat(
+                                    "defaultScores: " + scores1[j] + " bulkScores: " + scores2[j],
+                                    diff / scores2[j],
+                                    lessThan(1e-5f)
+                                );
                             } else {
                                 assertEquals(scores1[j], scores2[j], 1e-2f);
                             }

muted-tests.yml

@@ -577,6 +577,18 @@ tests:
 - class: org.elasticsearch.server.cli.MachineDependentHeapTests
   method: testMlOnlyOptions
   issue: https://github.com/elastic/elasticsearch/issues/129236
+- class: org.elasticsearch.upgrades.RunningSnapshotIT
+  method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=1}
+  issue: https://github.com/elastic/elasticsearch/issues/129644
+- class: org.elasticsearch.upgrades.RunningSnapshotIT
+  method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=2}
+  issue: https://github.com/elastic/elasticsearch/issues/129645
+- class: org.elasticsearch.upgrades.RunningSnapshotIT
+  method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=3}
+  issue: https://github.com/elastic/elasticsearch/issues/129646
+- class: org.elasticsearch.test.apmintegration.TracesApmIT
+  method: testApmIntegration
+  issue: https://github.com/elastic/elasticsearch/issues/129651
 
 # Examples:
 #

server/src/internalClusterTest/java/org/elasticsearch/search/query/VectorIT.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.search.query;
 
+import org.apache.lucene.util.VectorUtil;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
@@ -30,10 +31,11 @@ public class VectorIT extends ESIntegTestCase {
     private static final String VECTOR_FIELD = "vector";
     private static final String NUM_ID_FIELD = "num_id";
 
-    private static void randomVector(float[] vector) {
+    private static void randomVector(float[] vector, int constant) {
         for (int i = 0; i < vector.length; i++) {
-            vector[i] = randomFloat();
+            vector[i] = randomFloat() * constant;
         }
+        VectorUtil.l2normalize(vector);
     }
 
     @Before
@@ -43,6 +45,7 @@ public void setup() throws IOException {
             .startObject("properties")
             .startObject(VECTOR_FIELD)
             .field("type", "dense_vector")
+            .field("similarity", "dot_product")
             .startObject("index_options")
             .field("type", "hnsw")
             .endObject()
@@ -59,20 +62,21 @@ public void setup() throws IOException {
             .build();
         prepareCreate(INDEX_NAME).setMapping(mapping).setSettings(settings).get();
         ensureGreen(INDEX_NAME);
+        float[] vector = new float[16];
         for (int i = 0; i < 150; i++) {
-            float[] vector = new float[8];
-            randomVector(vector);
+            randomVector(vector, i % 25 + 1);
             prepareIndex(INDEX_NAME).setId(Integer.toString(i)).setSource(VECTOR_FIELD, vector, NUM_ID_FIELD, i).get();
         }
         forceMerge(true);
         refresh(INDEX_NAME);
     }
 
     public void testFilteredQueryStrategy() {
-        float[] vector = new float[8];
-        randomVector(vector);
+        float[] vector = new float[16];
+        randomVector(vector, 25);
+        int upperLimit = 35;
         var query = new KnnSearchBuilder(VECTOR_FIELD, vector, 1, 1, null, null).addFilterQuery(
-            QueryBuilders.rangeQuery(NUM_ID_FIELD).lte(30)
+            QueryBuilders.rangeQuery(NUM_ID_FIELD).lte(35)
         );
         assertResponse(client().prepareSearch(INDEX_NAME).setKnnSearch(List.of(query)).setSize(1).setProfile(true), acornResponse -> {
             assertNotEquals(0, acornResponse.getHits().getHits().length);
@@ -116,6 +120,8 @@ public void testFilteredQueryStrategy() {
                 assertTrue(
                     "fanoutVectorOps [" + fanoutVectorOpsSum + "] is not gt acornVectorOps [" + vectorOpsSum + "]",
                     fanoutVectorOpsSum > vectorOpsSum
+                        // if both switch to brute-force due to excessive exploration, they will both equal to upperLimit
+                        || (fanoutVectorOpsSum == vectorOpsSum && vectorOpsSum == upperLimit + 1)
                 );
             });
         });

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -198,7 +198,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion IDP_CUSTOM_SAML_ATTRIBUTES_ALLOW_LIST_8_19 = def(8_841_0_50);
     public static final TransportVersion SETTINGS_IN_DATA_STREAMS_8_19 = def(8_841_0_51);
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_REMOVE_ERROR_PARSING_8_19 = def(8_841_0_52);
-    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE_8_19 = def(8_841_0_53);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_EMBEDDING_BATCH_SIZE_8_19 = def(8_841_0_53);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE_8_19 = def(8_841_0_54);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);
@@ -305,7 +306,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion STATE_PARAM_GET_SNAPSHOT = def(9_100_0_00);
     public static final TransportVersion PROJECT_ID_IN_SNAPSHOTS_DELETIONS_AND_REPO_CLEANUP = def(9_101_0_00);
     public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_REMOVE_ERROR_PARSING = def(9_102_0_00);
-    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE = def(9_103_0_00);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_EMBEDDING_BATCH_SIZE = def(9_103_0_00);
+    public static final TransportVersion ML_INFERENCE_CUSTOM_SERVICE_INPUT_TYPE = def(9_104_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/inference/ChunkingStrategy.java

@@ -16,6 +16,7 @@
 public enum ChunkingStrategy {
     WORD("word"),
     SENTENCE("sentence"),
+    RECURSIVE("recursive"),
     NONE("none");
 
     private final String chunkingStrategy;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceNamedWriteablesProvider.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
 import org.elasticsearch.xpack.inference.action.task.StreamingTaskManager;
 import org.elasticsearch.xpack.inference.chunking.NoneChunkingSettings;
+import org.elasticsearch.xpack.inference.chunking.RecursiveChunkingSettings;
 import org.elasticsearch.xpack.inference.chunking.SentenceBoundaryChunkingSettings;
 import org.elasticsearch.xpack.inference.chunking.WordBoundaryChunkingSettings;
 import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
@@ -567,6 +568,9 @@ private static void addChunkingSettingsNamedWriteables(List<NamedWriteableRegist
                 SentenceBoundaryChunkingSettings::new
             )
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(ChunkingSettings.class, RecursiveChunkingSettings.NAME, RecursiveChunkingSettings::new)
+        );
     }
 
     private static void addInferenceResultsNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/ChunkerBuilder.java

@@ -19,6 +19,7 @@ public static Chunker fromChunkingStrategy(ChunkingStrategy chunkingStrategy) {
             case NONE -> NoopChunker.INSTANCE;
             case WORD -> new WordBoundaryChunker();
             case SENTENCE -> new SentenceBoundaryChunker();
+            case RECURSIVE -> new RecursiveChunker();
         };
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/ChunkerUtils.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.chunking;
+
+import com.ibm.icu.text.BreakIterator;
+
+public class ChunkerUtils {
+
+    // setText() should be applied before using this function.
+    static int countWords(int start, int end, BreakIterator wordIterator) {
+        assert start < end;
+        wordIterator.preceding(start); // start of the current word
+
+        int boundary = wordIterator.current();
+        int wordCount = 0;
+        while (boundary != BreakIterator.DONE && boundary <= end) {
+            int wordStatus = wordIterator.getRuleStatus();
+            if (wordStatus != BreakIterator.WORD_NONE) {
+                wordCount++;
+            }
+            boundary = wordIterator.next();
+        }
+
+        return wordCount;
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/ChunkingSettingsBuilder.java

@@ -48,6 +48,7 @@ public static ChunkingSettings fromMap(Map<String, Object> settings, boolean ret
             case NONE -> NoneChunkingSettings.INSTANCE;
             case WORD -> WordBoundaryChunkingSettings.fromMap(new HashMap<>(settings));
             case SENTENCE -> SentenceBoundaryChunkingSettings.fromMap(new HashMap<>(settings));
+            case RECURSIVE -> RecursiveChunkingSettings.fromMap(new HashMap<>(settings));
         };
     }
 }