Merge branch 'main' into lucene_10_2_1

Author: ChrisHegarty

docs/changelog/127225.yaml

@@ -0,0 +1,6 @@
+pr: 127225
+summary: Fix count optimization with pushable union types
+area: ES|QL
+type: bug
+issues:
+ - 127200

muted-tests.yml

@@ -408,9 +408,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.multi_node.EsqlSpecIT
   method: test {rerank.Reranker before a limit ASYNC}
   issue: https://github.com/elastic/elasticsearch/issues/127051
-- class: org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecIT
-  method: test {rrf.SimpleRrf ASYNC}
-  issue: https://github.com/elastic/elasticsearch/issues/127063
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test026InstallBundledRepositoryPlugins
   issue: https://github.com/elastic/elasticsearch/issues/127081
@@ -426,9 +423,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.single_node.GenerativeIT
   method: test
   issue: https://github.com/elastic/elasticsearch/issues/127157
-- class: org.elasticsearch.xpack.esql.qa.single_node.EsqlSpecIT
-  method: test {fork.ForkWithWhereSortDescAndLimit SYNC}
-  issue: https://github.com/elastic/elasticsearch/issues/127326
 - class: org.elasticsearch.geometry.utils.SpatialEnvelopeVisitorTests
   method: testVisitGeoPointsWrapping
   issue: https://github.com/elastic/elasticsearch/issues/123425
@@ -450,6 +444,9 @@ tests:
 - class: org.elasticsearch.action.admin.cluster.state.TransportClusterStateActionDisruptionIT
   method: testLocalRequestAlwaysSucceeds
   issue: https://github.com/elastic/elasticsearch/issues/127423
+- class: org.elasticsearch.action.admin.cluster.state.TransportClusterStateActionDisruptionIT
+  method: testLocalRequestWaitsForMetadata
+  issue: https://github.com/elastic/elasticsearch/issues/127466
 
 # Examples:
 #

server/src/main/java/org/elasticsearch/index/codec/vectors/KMeans.java

@@ -0,0 +1,115 @@
+/*
+ * @notice
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Modifications copyright (C) 2025 Elasticsearch B.V.
+ */
+package org.elasticsearch.index.codec.vectors;
+
+import org.apache.lucene.codecs.lucene95.HasIndexSlice;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.util.Bits;
+
+import java.io.IOException;
+import java.util.Random;
+import java.util.function.IntUnaryOperator;
+
+class SampleReader extends FloatVectorValues implements HasIndexSlice {
+    private final FloatVectorValues origin;
+    private final int sampleSize;
+    private final IntUnaryOperator sampleFunction;
+
+    SampleReader(FloatVectorValues origin, int sampleSize, IntUnaryOperator sampleFunction) {
+        this.origin = origin;
+        this.sampleSize = sampleSize;
+        this.sampleFunction = sampleFunction;
+    }
+
+    @Override
+    public int size() {
+        return sampleSize;
+    }
+
+    @Override
+    public int dimension() {
+        return origin.dimension();
+    }
+
+    @Override
+    public FloatVectorValues copy() throws IOException {
+        throw new IllegalStateException("Not supported");
+    }
+
+    @Override
+    public IndexInput getSlice() {
+        return ((HasIndexSlice) origin).getSlice();
+    }
+
+    @Override
+    public float[] vectorValue(int targetOrd) throws IOException {
+        return origin.vectorValue(sampleFunction.applyAsInt(targetOrd));
+    }
+
+    @Override
+    public int getVectorByteLength() {
+        return origin.getVectorByteLength();
+    }
+
+    @Override
+    public int ordToDoc(int ord) {
+        throw new IllegalStateException("Not supported");
+    }
+
+    @Override
+    public Bits getAcceptOrds(Bits acceptDocs) {
+        throw new IllegalStateException("Not supported");
+    }
+
+    static SampleReader createSampleReader(FloatVectorValues origin, int k, long seed) {
+        // TODO can we do something algorithmically that aligns an ordinal with a unique integer between 0 and numVectors?
+        if (k >= origin.size()) {
+            new SampleReader(origin, origin.size(), i -> i);
+        }
+        // TODO maybe use bigArrays?
+        int[] samples = reservoirSample(origin.size(), k, seed);
+        return new SampleReader(origin, samples.length, i -> samples[i]);
+    }
+
+    /**
+     * Sample k elements from n elements according to reservoir sampling algorithm.
+     *
+     * @param n number of elements
+     * @param k number of samples
+     * @param seed random seed
+     * @return array of k samples
+     */
+    public static int[] reservoirSample(int n, int k, long seed) {
+        Random rnd = new Random(seed);
+        int[] reservoir = new int[k];
+        for (int i = 0; i < k; i++) {
+            reservoir[i] = i;
+        }
+        for (int i = k; i < n; i++) {
+            int j = rnd.nextInt(i + 1);
+            if (j < k) {
+                reservoir[j] = i;
+            }
+        }
+        return reservoir;
+    }
+
+}

server/src/main/java/org/elasticsearch/index/codec/vectors/SampleReader.java

@@ -23,6 +23,7 @@
 import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.hnsw.FlatVectorsWriter;
+import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsReader;
 import org.apache.lucene.codecs.lucene99.Lucene99FlatVectorsWriter;
 import org.apache.lucene.index.SegmentReadState;
 import org.apache.lucene.index.SegmentWriteState;
@@ -62,7 +63,16 @@ public FlatVectorsWriter fieldsWriter(SegmentWriteState state) throws IOExceptio
 
     @Override
     public FlatVectorsReader fieldsReader(SegmentReadState state) throws IOException {
-        return new DirectIOLucene99FlatVectorsReader(state, vectorsScorer);
+        if (DirectIOLucene99FlatVectorsReader.shouldUseDirectIO(state)) {
+            // Use mmap for merges and direct I/O for searches.
+            // TODO: Open the mmap file with sequential access instead of random (current behavior).
+            return new MergeReaderWrapper(
+                new DirectIOLucene99FlatVectorsReader(state, vectorsScorer),
+                new Lucene99FlatVectorsReader(state, vectorsScorer)
+            );
+        } else {
+            return new Lucene99FlatVectorsReader(state, vectorsScorer);
+        }
     }
 
     @Override

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsFormat.java

@@ -85,6 +85,10 @@ public DirectIOLucene99FlatVectorsReader(SegmentReadState state, FlatVectorsScor
         }
     }
 
+    public static boolean shouldUseDirectIO(SegmentReadState state) {
+        return USE_DIRECT_IO && state.directory instanceof DirectIOIndexInputSupplier;
+    }
+
     private int readMetadata(SegmentReadState state) throws IOException {
         String metaFileName = IndexFileNames.segmentFileName(
             state.segmentInfo.name,
@@ -122,11 +126,9 @@ private static IndexInput openDataInput(
     ) throws IOException {
         String fileName = IndexFileNames.segmentFileName(state.segmentInfo.name, state.segmentSuffix, fileExtension);
         // use direct IO for accessing raw vector data for searches
-        IndexInput in = USE_DIRECT_IO
-            && context.context() == IOContext.Context.DEFAULT
-            && state.directory instanceof DirectIOIndexInputSupplier did
-                ? did.openInputDirect(fileName, context)
-                : state.directory.openInput(fileName, context);
+        IndexInput in = USE_DIRECT_IO && state.directory instanceof DirectIOIndexInputSupplier did
+            ? did.openInputDirect(fileName, context)
+            : state.directory.openInput(fileName, context);
         boolean success = false;
         try {
             int versionVectorData = CodecUtil.checkIndexHeader(

server/src/main/java/org/elasticsearch/index/codec/vectors/es818/DirectIOLucene99FlatVectorsReader.java

@@ -0,0 +1,89 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.codec.vectors.es818;
+
+import org.apache.lucene.codecs.hnsw.FlatVectorsReader;
+import org.apache.lucene.index.ByteVectorValues;
+import org.apache.lucene.index.FloatVectorValues;
+import org.apache.lucene.search.KnnCollector;
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.hnsw.RandomVectorScorer;
+import org.elasticsearch.core.IOUtils;
+
+import java.io.IOException;
+import java.util.Collection;
+
+class MergeReaderWrapper extends FlatVectorsReader {
+
+    private final FlatVectorsReader mainReader;
+    private final FlatVectorsReader mergeReader;
+
+    protected MergeReaderWrapper(FlatVectorsReader mainReader, FlatVectorsReader mergeReader) {
+        super(mainReader.getFlatVectorScorer());
+        this.mainReader = mainReader;
+        this.mergeReader = mergeReader;
+    }
+
+    @Override
+    public RandomVectorScorer getRandomVectorScorer(String field, float[] target) throws IOException {
+        return mainReader.getRandomVectorScorer(field, target);
+    }
+
+    @Override
+    public RandomVectorScorer getRandomVectorScorer(String field, byte[] target) throws IOException {
+        return mainReader.getRandomVectorScorer(field, target);
+    }
+
+    @Override
+    public void checkIntegrity() throws IOException {
+        mainReader.checkIntegrity();
+    }
+
+    @Override
+    public FloatVectorValues getFloatVectorValues(String field) throws IOException {
+        return mainReader.getFloatVectorValues(field);
+    }
+
+    @Override
+    public ByteVectorValues getByteVectorValues(String field) throws IOException {
+        return mainReader.getByteVectorValues(field);
+    }
+
+    @Override
+    public FlatVectorsReader getMergeInstance() {
+        return mergeReader;
+    }
+
+    @Override
+    public long ramBytesUsed() {
+        return mainReader.ramBytesUsed();
+    }
+
+    @Override
+    public Collection<Accountable> getChildResources() {
+        return mainReader.getChildResources();
+    }
+
+    @Override
+    public void search(String field, float[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
+        mainReader.search(field, target, knnCollector, acceptDocs);
+    }
+
+    @Override
+    public void search(String field, byte[] target, KnnCollector knnCollector, Bits acceptDocs) throws IOException {
+        mainReader.search(field, target, knnCollector, acceptDocs);
+    }
+
+    @Override
+    public void close() throws IOException {
+        IOUtils.close(mainReader, mergeReader);
+    }
+}