Merge branch 'main' into wire_cps_check

Author: idegtiarenko

docs/changelog/137220.yaml

@@ -0,0 +1,5 @@
+pr: 137220
+summary: Skip dataframes when disabled
+area: Machine Learning
+type: bug
+issues: []

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSBitToInt4ESNextOSQVectorsScorer.java

@@ -14,7 +14,6 @@
 import jdk.incubator.vector.LongVector;
 import jdk.incubator.vector.ShortVector;
 import jdk.incubator.vector.VectorOperators;
-import jdk.incubator.vector.VectorSpecies;
 
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.IndexInput;
@@ -31,23 +30,6 @@
 /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */
 final class MSBitToInt4ESNextOSQVectorsScorer extends MemorySegmentESNextOSQVectorsScorer.MemorySegmentScorer {
 
-    private static final int BULK_SIZE = MemorySegmentESNextOSQVectorsScorer.BULK_SIZE;
-    private static final float FOUR_BIT_SCALE = 1f / ((1 << 4) - 1);
-
-    private static final VectorSpecies<Integer> INT_SPECIES_128 = IntVector.SPECIES_128;
-
-    private static final VectorSpecies<Long> LONG_SPECIES_128 = LongVector.SPECIES_128;
-    private static final VectorSpecies<Long> LONG_SPECIES_256 = LongVector.SPECIES_256;
-
-    private static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;
-    private static final VectorSpecies<Byte> BYTE_SPECIES_256 = ByteVector.SPECIES_256;
-
-    private static final VectorSpecies<Short> SHORT_SPECIES_128 = ShortVector.SPECIES_128;
-    private static final VectorSpecies<Short> SHORT_SPECIES_256 = ShortVector.SPECIES_256;
-
-    private static final VectorSpecies<Float> FLOAT_SPECIES_128 = FloatVector.SPECIES_128;
-    private static final VectorSpecies<Float> FLOAT_SPECIES_256 = FloatVector.SPECIES_256;
-
     MSBitToInt4ESNextOSQVectorsScorer(IndexInput in, int dimensions, int dataLength, MemorySegment memorySegment) {
         super(in, dimensions, dataLength, memorySegment);
     }

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MSInt4SymmetricESNextOSQVectorsScorer.java

@@ -25,20 +25,6 @@
 /** Panamized scorer for quantized vectors stored as a {@link MemorySegment}. */
 public final class MemorySegmentESNextOSQVectorsScorer extends ESNextOSQVectorsScorer {
 
-    private static final VectorSpecies<Integer> INT_SPECIES_128 = IntVector.SPECIES_128;
-
-    private static final VectorSpecies<Long> LONG_SPECIES_128 = LongVector.SPECIES_128;
-    private static final VectorSpecies<Long> LONG_SPECIES_256 = LongVector.SPECIES_256;
-
-    private static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;
-    private static final VectorSpecies<Byte> BYTE_SPECIES_256 = ByteVector.SPECIES_256;
-
-    private static final VectorSpecies<Short> SHORT_SPECIES_128 = ShortVector.SPECIES_128;
-    private static final VectorSpecies<Short> SHORT_SPECIES_256 = ShortVector.SPECIES_256;
-
-    private static final VectorSpecies<Float> FLOAT_SPECIES_128 = FloatVector.SPECIES_128;
-    private static final VectorSpecies<Float> FLOAT_SPECIES_256 = FloatVector.SPECIES_256;
-
     private final MemorySegment memorySegment;
     private final MemorySegmentScorer scorer;
 
@@ -54,6 +40,10 @@ public MemorySegmentESNextOSQVectorsScorer(
         this.memorySegment = memorySegment;
         if (queryBits == 4 && indexBits == 1) {
             this.scorer = new MSBitToInt4ESNextOSQVectorsScorer(in, dimensions, dataLength, memorySegment);
+        } else if (queryBits == 4 && indexBits == 4) {
+            this.scorer = new MSInt4SymmetricESNextOSQVectorsScorer(in, dimensions, dataLength, memorySegment);
+        } else if (queryBits == 4 && indexBits == 2) {
+            throw new IllegalArgumentException("Only symmetric 4-bit query and 1-bit index supported");
         } else {
             throw new IllegalArgumentException("Only asymmetric 4-bit query and 1-bit index supported");
         }
@@ -112,7 +102,24 @@ public float scoreBulk(
         );
     }
 
-    abstract static sealed class MemorySegmentScorer permits MSBitToInt4ESNextOSQVectorsScorer {
+    abstract static sealed class MemorySegmentScorer permits MSBitToInt4ESNextOSQVectorsScorer, MSInt4SymmetricESNextOSQVectorsScorer {
+
+        static final int BULK_SIZE = MemorySegmentESNextOSQVectorsScorer.BULK_SIZE;
+        static final float FOUR_BIT_SCALE = 1f / ((1 << 4) - 1);
+        static final VectorSpecies<Integer> INT_SPECIES_128 = IntVector.SPECIES_128;
+
+        static final VectorSpecies<Long> LONG_SPECIES_128 = LongVector.SPECIES_128;
+        static final VectorSpecies<Long> LONG_SPECIES_256 = LongVector.SPECIES_256;
+
+        static final VectorSpecies<Byte> BYTE_SPECIES_128 = ByteVector.SPECIES_128;
+        static final VectorSpecies<Byte> BYTE_SPECIES_256 = ByteVector.SPECIES_256;
+
+        static final VectorSpecies<Short> SHORT_SPECIES_128 = ShortVector.SPECIES_128;
+        static final VectorSpecies<Short> SHORT_SPECIES_256 = ShortVector.SPECIES_256;
+
+        static final VectorSpecies<Float> FLOAT_SPECIES_128 = FloatVector.SPECIES_128;
+        static final VectorSpecies<Float> FLOAT_SPECIES_256 = FloatVector.SPECIES_256;
+
         protected final MemorySegment memorySegment;
         protected final IndexInput in;
         protected final int length;

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/MemorySegmentESNextOSQVectorsScorer.java

@@ -40,7 +40,7 @@ public ESNextOSQVectorsScorer newESNextOSQVectorsScorer(IndexInput input, byte q
         if (PanamaESVectorUtilSupport.HAS_FAST_INTEGER_VECTORS
             && input instanceof MemorySegmentAccessInput msai
             && queryBits == 4
-            && indexBits == 1) {
+            && (indexBits == 1 || indexBits == 4)) {
             MemorySegment ms = msai.segmentSliceOrNull(0, input.length());
             if (ms != null) {
                 return new MemorySegmentESNextOSQVectorsScorer(input, queryBits, indexBits, dimension, dataLength, ms);

libs/simdvec/src/main21/java/org/elasticsearch/simdvec/internal/vectorization/PanamaESVectorizationProvider.java

@@ -492,9 +492,6 @@ tests:
 - class: org.elasticsearch.readiness.ReadinessClusterIT
   method: testReadinessDuringRestartsNormalOrder
   issue: https://github.com/elastic/elasticsearch/issues/136955
-- class: org.elasticsearch.xpack.ilm.TimeSeriesDataStreamsIT
-  method: testSearchableSnapshotAction
-  issue: https://github.com/elastic/elasticsearch/issues/137167
 - class: org.elasticsearch.xpack.security.CoreWithSecurityClientYamlTestSuiteIT
   method: test {yaml=indices.validate_query/20_query_string/validate_query with query_string parameters}
   issue: https://github.com/elastic/elasticsearch/issues/137391

muted-tests.yml

@@ -17,7 +17,6 @@
 import org.apache.http.client.CredentialsProvider;
 import org.apache.http.impl.client.BasicCredentialsProvider;
 import org.apache.logging.log4j.core.config.plugins.util.PluginManager;
-import org.apache.lucene.util.IOConsumer;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
@@ -31,7 +30,6 @@
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
-import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xcontent.XContentType;
 
@@ -45,7 +43,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.Semaphore;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -58,7 +55,6 @@
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.reader;
 
 public class CsvTestsDataLoader {
-    private static final int PARALLEL_THREADS = 10;
     private static final int BULK_DATA_SIZE = 100_000;
     private static final TestDataset EMPLOYEES = new TestDataset("employees", "mapping-default.json", "employees.csv").noSubfields();
     private static final TestDataset EMPLOYEES_INCOMPATIBLE = new TestDataset(
@@ -433,42 +429,18 @@ private static void loadDataSetIntoEs(
         IndexCreator indexCreator
     ) throws IOException {
         Logger logger = LogManager.getLogger(CsvTestsDataLoader.class);
-        List<TestDataset> datasets = availableDatasetsForEs(
-            supportsIndexModeLookup,
-            supportsSourceFieldMapping,
-            inferenceEnabled,
-            timeSeriesOnly
-        ).stream().toList();
-
-        logger.info("Creating test indices");
-        executeInParallel(datasets, dataset -> createIndex(client, dataset, indexCreator), "Failed to create indices in parallel");
 
+        Set<String> loadedDatasets = new HashSet<>();
         logger.info("Loading test datasets");
-        executeInParallel(datasets, dataset -> loadData(client, dataset, logger), "Failed to load data in parallel");
-
-        forceMerge(client, datasets.stream().map(d -> d.indexName).collect(Collectors.toSet()), logger);
-
+        for (var dataset : availableDatasetsForEs(supportsIndexModeLookup, supportsSourceFieldMapping, inferenceEnabled, timeSeriesOnly)) {
+            load(client, dataset, logger, indexCreator);
+            loadedDatasets.add(dataset.indexName);
+        }
+        forceMerge(client, loadedDatasets, logger);
         logger.info("Loading enrich policies");
-        executeInParallel(
-            ENRICH_POLICIES,
-            policy -> loadEnrichPolicy(client, policy.policyName, policy.policyFileName, logger),
-            "Failed to load enrich policies in parallel"
-        );
-
-    }
-
-    private static <T> void executeInParallel(List<T> items, IOConsumer<T> consumer, String errorMessage) {
-        Semaphore semaphore = new Semaphore(PARALLEL_THREADS);
-        ESTestCase.runInParallel(items.size(), i -> {
-            try {
-                semaphore.acquire();
-                consumer.accept(items.get(i));
-            } catch (IOException | InterruptedException e) {
-                throw new RuntimeException(errorMessage, e);
-            } finally {
-                semaphore.release();
-            }
-        });
+        for (var policy : ENRICH_POLICIES) {
+            loadEnrichPolicy(client, policy.policyName, policy.policyFileName, logger);
+        }
     }
 
     public static void createInferenceEndpoints(RestClient client) throws IOException {
@@ -626,14 +598,12 @@ private static URL getResource(String name) {
         return result;
     }
 
-    private static void createIndex(RestClient client, TestDataset dataset, IndexCreator indexCreator) throws IOException {
+    private static void load(RestClient client, TestDataset dataset, Logger logger, IndexCreator indexCreator) throws IOException {
+        logger.info("Loading dataset [{}] into ES index [{}]", dataset.dataFileName, dataset.indexName);
         URL mapping = getResource("/" + dataset.mappingFileName);
         Settings indexSettings = dataset.readSettingsFile();
         indexCreator.createIndex(client, dataset.indexName, readMappingFile(mapping, dataset.typeMapping), indexSettings);
-    }
 
-    private static void loadData(RestClient client, TestDataset dataset, Logger logger) throws IOException {
-        logger.info("Loading dataset [{}] into ES index [{}]", dataset.dataFileName, dataset.indexName);
         // Some examples only test that the query and mappings are valid, and don't need example data. Use .noData() for those
         if (dataset.dataFileName != null) {
             URL data = getResource("/data/" + dataset.dataFileName);

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvTestsDataLoader.java

@@ -11,18 +11,12 @@
 
 import java.net.ConnectException;
 
-import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.startsWith;
 
 public class CsvTestsDataLoaderTests extends ESTestCase {
 
     public void testCsvTestsDataLoaderExecution() {
-        Throwable cause = expectThrows(AssertionError.class, () -> CsvTestsDataLoader.main(new String[] {}));
-        // find the root cause
-        while (cause.getCause() != null) {
-            cause = cause.getCause();
-        }
-        assertThat(cause, instanceOf(ConnectException.class));
-        assertThat(cause.getMessage(), startsWith("Connection refused"));
+        ConnectException ce = expectThrows(ConnectException.class, () -> CsvTestsDataLoader.main(new String[] {}));
+        assertThat(ce.getMessage(), startsWith("Connection refused"));
     }
 }

x-pack/plugin/esql/qa/testFixtures/src/test/java/org/elasticsearch/xpack/esql/CsvTestsDataLoaderTests.java

@@ -184,8 +184,8 @@ public void testSearchableSnapshotAction() throws Exception {
         // Manual rollover the original index such that it's not the write index in the data stream anymore
         rolloverMaxOneDocCondition(client(), dataStream);
 
-        awaitIndexExists(restoredIndexName);
-        awaitIndexDoesNotExist(backingIndexName, TimeValue.timeValueSeconds(60));
+        awaitIndexExists(restoredIndexName, TimeValue.timeValueSeconds(20));
+        awaitIndexDoesNotExist(backingIndexName);
         assertBusy(
             () -> assertThat(explainIndex(client(), restoredIndexName).get("step"), is(PhaseCompleteStep.NAME)),
             30,