Merge branch 'main' into retry_shard_movements_during_query

# Conflicts:
#	x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java
#	x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSender.java
#	x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSenderTests.java

Author: idegtiarenko

docs/changelog/125631.yaml

@@ -0,0 +1,5 @@
+pr: 125631
+summary: Add `documents_found` and `values_loaded`
+area: ES|QL
+type: enhancement
+issues: []

docs/changelog/126911.yaml

@@ -0,0 +1,6 @@
+pr: 126911
+summary: Fix `vec_caps` to test for OS support too (on x64)
+area: Vector Search
+type: bug
+issues:
+ - 126809

libs/native/libraries/build.gradle

@@ -19,7 +19,7 @@ configurations {
 }
 
 var zstdVersion = "1.5.5"
-var vecVersion = "1.0.10"
+var vecVersion = "1.0.11"
 
 repositories {
   exclusiveContent {

libs/native/src/main/java/org/elasticsearch/nativeaccess/jdk/JdkVectorLibrary.java

@@ -41,7 +41,7 @@ public final class JdkVectorLibrary implements VectorLibrary {
         try {
             int caps = (int) vecCaps$mh.invokeExact();
             logger.info("vec_caps=" + caps);
-            if (caps != 0) {
+            if (caps > 0) {
                 if (caps == 2) {
                     dot7u$mh = downcallHandle(
                         "dot7u_2",
@@ -67,6 +67,11 @@ public final class JdkVectorLibrary implements VectorLibrary {
                 }
                 INSTANCE = new JdkVectorSimilarityFunctions();
             } else {
+                if (caps < 0) {
+                    logger.warn("""
+                        Your CPU supports vector capabilities, but they are disabled at OS level. For optimal performance, \
+                        enable them in your OS/Hypervisor/VM/container""");
+                }
                 dot7u$mh = null;
                 sqr7u$mh = null;
                 INSTANCE = null;

libs/simdvec/native/build.gradle

@@ -11,7 +11,7 @@ apply plugin: 'cpp'
 
 var os = org.gradle.internal.os.OperatingSystem.current()
 
-// To update this library run publish_vec_binaries.sh  ( or ./gradlew vecSharedLibrary )
+// To update this library run publish_vec_binaries.sh  ( or ./gradlew buildSharedLibrary )
 // Or
 // For local development, build the docker image with:
 //   docker build --platform linux/arm64 --progress=plain --file=Dockerfile.aarch64 . (for aarch64)

libs/simdvec/native/publish_vec_binaries.sh

@@ -20,7 +20,7 @@ if [ -z "$ARTIFACTORY_API_KEY" ]; then
   exit 1;
 fi
 
-VERSION="1.0.10"
+VERSION="1.0.11"
 ARTIFACTORY_REPOSITORY="${ARTIFACTORY_REPOSITORY:-https://artifactory.elastic.dev/artifactory/elasticsearch-native/}"
 TEMP=$(mktemp -d)
 

libs/simdvec/native/src/vec/c/amd64/vec.c

@@ -46,6 +46,23 @@ static inline void cpuid(int output[4], int functionNumber) {
 #endif
 }
 
+// Multi-platform XGETBV "intrinsic"
+static inline int64_t xgetbv(int ctr) {
+#if defined(__GNUC__) || defined(__clang__)
+    // use inline assembly, Gnu/AT&T syntax
+    uint32_t a, d;
+    __asm("xgetbv" : "=a"(a),"=d"(d) : "c"(ctr) : );
+    return a | (((uint64_t) d) << 32);
+
+#elif (defined (_MSC_FULL_VER) && _MSC_FULL_VER >= 160040000) || (defined (__INTEL_COMPILER) && __INTEL_COMPILER >= 1200)
+    // Microsoft or Intel compiler supporting _xgetbv intrinsic
+    return _xgetbv(ctr);
+
+#else
+   #error Unsupported compiler
+#endif
+}
+
 // Utility function to horizontally add 8 32-bit integers
 static inline int hsum_i32_8(const __m256i a) {
     const __m128i sum128 = _mm_add_epi32(_mm256_castsi256_si128(a), _mm256_extractf128_si256(a, 1));
@@ -57,11 +74,20 @@ static inline int hsum_i32_8(const __m256i a) {
 
 EXPORT int vec_caps() {
     int cpuInfo[4] = {-1};
-    // Calling __cpuid with 0x0 as the function_id argument
+    // Calling CPUID function 0x0 as the function_id argument
     // gets the number of the highest valid function ID.
     cpuid(cpuInfo, 0);
     int functionIds = cpuInfo[0];
+    if (functionIds == 0) {
+        // No CPUID functions
+        return 0;
+    }
+    // call CPUID function 0x1 for feature flags
+    cpuid(cpuInfo, 1);
+    int hasOsXsave = (cpuInfo[2] & (1 << 27)) != 0;
+    int avxEnabledInOS = hasOsXsave && ((xgetbv(0) & 6) == 6);
     if (functionIds >= 7) {
+        // call CPUID function 0x7 for AVX2/512 flags
         cpuid(cpuInfo, 7);
         int ebx = cpuInfo[1];
         int ecx = cpuInfo[2];
@@ -72,10 +98,18 @@ EXPORT int vec_caps() {
         // int avx512_vnni = (ecx & 0x00000800) != 0;
         // if (avx512 && avx512_vnni) {
         if (avx512) {
-            return 2;
+            if (avxEnabledInOS) {
+                return 2;
+            } else {
+                return -2;
+            }
         }
         if (avx2) {
-            return 1;
+            if (avxEnabledInOS) {
+                return 1;
+            } else {
+                return -1;
+            }
         }
     }
     return 0;

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

@@ -222,6 +222,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion AMAZON_BEDROCK_TASK_SETTINGS = def(9_049_0_00);
     public static final TransportVersion ESQL_REPORT_SHARD_PARTITIONING = def(9_050_0_00);
     public static final TransportVersion ESQL_QUERY_PLANNING_DURATION = def(9_051_0_00);
+    public static final TransportVersion ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED = def(9_052_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/common/Strings.java

@@ -822,7 +822,7 @@ public static String toString(ChunkedToXContent chunkedToXContent, boolean prett
      * Allows to configure the params.
      * Allows to control whether the outputted json needs to be pretty printed and human readable.
      */
-    private static String toString(ToXContent toXContent, ToXContent.Params params, boolean pretty, boolean human) {
+    public static String toString(ToXContent toXContent, ToXContent.Params params, boolean pretty, boolean human) {
         try {
             XContentBuilder builder = createBuilder(pretty, human);
             if (toXContent.isFragment()) {

test/framework/src/main/java/org/elasticsearch/node/MockNode.java

@@ -9,6 +9,8 @@
 
 package org.elasticsearch.node;
 
+import org.elasticsearch.action.search.OnlinePrewarmingService;
+import org.elasticsearch.action.search.OnlinePrewarmingServiceProvider;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.ClusterInfoService;
 import org.elasticsearch.cluster.MockInternalClusterInfoService;
@@ -117,6 +119,10 @@ SearchService newSearchService(
                     tracer
                 );
             }
+            OnlinePrewarmingService onlinePrewarmingService = pluginsService.loadSingletonServiceProvider(
+                OnlinePrewarmingServiceProvider.class,
+                () -> OnlinePrewarmingServiceProvider.DEFAULT
+            ).create(clusterService.getSettings(), threadPool, clusterService);
             return new MockSearchService(
                 clusterService,
                 indicesService,
@@ -126,7 +132,8 @@ SearchService newSearchService(
                 fetchPhase,
                 circuitBreakerService,
                 executorSelector,
-                tracer
+                tracer,
+                onlinePrewarmingService
             );
         }
 

test/framework/src/main/java/org/elasticsearch/search/MockSearchService.java

@@ -85,7 +85,8 @@ public MockSearchService(
         FetchPhase fetchPhase,
         CircuitBreakerService circuitBreakerService,
         ExecutorSelector executorSelector,
-        Tracer tracer
+        Tracer tracer,
+        OnlinePrewarmingService onlinePrewarmingService
     ) {
         super(
             clusterService,
@@ -97,7 +98,7 @@ public MockSearchService(
             circuitBreakerService,
             executorSelector,
             tracer,
-            OnlinePrewarmingService.NOOP
+            onlinePrewarmingService
         );
     }
 

test/framework/src/main/java/org/elasticsearch/test/rest/ESRestTestCase.java

@@ -2672,8 +2672,13 @@ protected static MapMatcher getProfileMatcher() {
             .entry("drivers", instanceOf(List.class));
     }
 
-    protected static MapMatcher getResultMatcher(boolean includeMetadata, boolean includePartial) {
+    protected static MapMatcher getResultMatcher(boolean includeMetadata, boolean includePartial, boolean includeDocumentsFound) {
         MapMatcher mapMatcher = matchesMap();
+        if (includeDocumentsFound) {
+            // Older versions may not return documents_found and values_loaded.
+            mapMatcher = mapMatcher.entry("documents_found", greaterThanOrEqualTo(0));
+            mapMatcher = mapMatcher.entry("values_loaded", greaterThanOrEqualTo(0));
+        }
         if (includeMetadata) {
             mapMatcher = mapMatcher.entry("took", greaterThanOrEqualTo(0));
         }
@@ -2688,7 +2693,7 @@ protected static MapMatcher getResultMatcher(boolean includeMetadata, boolean in
      * Create empty result matcher from result, taking into account all metadata items.
      */
     protected static MapMatcher getResultMatcher(Map<String, Object> result) {
-        return getResultMatcher(result.containsKey("took"), result.containsKey("is_partial"));
+        return getResultMatcher(result.containsKey("took"), result.containsKey("is_partial"), result.containsKey("documents_found"));
     }
 
     /**

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/CompositeBlock.java

@@ -83,7 +83,11 @@ public int getPositionCount() {
 
     @Override
     public int getTotalValueCount() {
-        throw new UnsupportedOperationException("Composite block");
+        int totalValueCount = 0;
+        for (Block b : blocks) {
+            totalValueCount += b.getTotalValueCount();
+        }
+        return totalValueCount;
     }
 
     @Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/LuceneOperator.java

@@ -434,6 +434,11 @@ public Map<String, LuceneSliceQueue.PartitioningStrategy> partitioningStrategies
             return partitioningStrategies;
         }
 
+        @Override
+        public long documentsFound() {
+            return rowsEmitted;
+        }
+
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/lucene/ValuesSourceReaderOperator.java

@@ -47,6 +47,8 @@
 import java.util.function.IntFunction;
 import java.util.function.Supplier;
 
+import static org.elasticsearch.TransportVersions.ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED;
+
 /**
  * Operator that extracts doc_values from a Lucene index out of pages that have been produced by {@link LuceneSourceOperator}
  * and outputs them to a new column.
@@ -113,6 +115,7 @@ public record ShardContext(IndexReader reader, Supplier<SourceLoader> newSourceL
     private final BlockFactory blockFactory;
 
     private final Map<String, Integer> readersBuilt = new TreeMap<>();
+    private long valuesLoaded;
 
     int lastShard = -1;
     int lastSegment = -1;
@@ -165,6 +168,9 @@ public int get(int i) {
                 }
             }
             success = true;
+            for (Block b : blocks) {
+                valuesLoaded += b.getTotalValueCount();
+            }
             return page.appendBlocks(blocks);
         } catch (IOException e) {
             throw new UncheckedIOException(e);
@@ -547,7 +553,7 @@ public String toString() {
 
     @Override
     protected Status status(long processNanos, int pagesProcessed, long rowsReceived, long rowsEmitted) {
-        return new Status(new TreeMap<>(readersBuilt), processNanos, pagesProcessed, rowsReceived, rowsEmitted);
+        return new Status(new TreeMap<>(readersBuilt), processNanos, pagesProcessed, rowsReceived, rowsEmitted, valuesLoaded);
     }
 
     public static class Status extends AbstractPageMappingOperator.Status {
@@ -558,21 +564,34 @@ public static class Status extends AbstractPageMappingOperator.Status {
         );
 
         private final Map<String, Integer> readersBuilt;
-
-        Status(Map<String, Integer> readersBuilt, long processNanos, int pagesProcessed, long rowsReceived, long rowsEmitted) {
+        private final long valuesLoaded;
+
+        Status(
+            Map<String, Integer> readersBuilt,
+            long processNanos,
+            int pagesProcessed,
+            long rowsReceived,
+            long rowsEmitted,
+            long valuesLoaded
+        ) {
             super(processNanos, pagesProcessed, rowsReceived, rowsEmitted);
             this.readersBuilt = readersBuilt;
+            this.valuesLoaded = valuesLoaded;
         }
 
         Status(StreamInput in) throws IOException {
             super(in);
             readersBuilt = in.readOrderedMap(StreamInput::readString, StreamInput::readVInt);
+            valuesLoaded = in.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED) ? in.readVLong() : 0;
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             super.writeTo(out);
             out.writeMap(readersBuilt, StreamOutput::writeVInt);
+            if (out.getTransportVersion().onOrAfter(ESQL_DOCUMENTS_FOUND_AND_VALUES_LOADED)) {
+                out.writeVLong(valuesLoaded);
+            }
         }
 
         @Override
@@ -584,6 +603,11 @@ public Map<String, Integer> readersBuilt() {
             return readersBuilt;
         }
 
+        @Override
+        public long valuesLoaded() {
+            return valuesLoaded;
+        }
+
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
@@ -592,6 +616,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
                 builder.field(e.getKey(), e.getValue());
             }
             builder.endObject();
+            builder.field("values_loaded", valuesLoaded);
             innerToXContent(builder);
             return builder.endObject();
         }
@@ -600,12 +625,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         public boolean equals(Object o) {
             if (super.equals(o) == false) return false;
             Status status = (Status) o;
-            return readersBuilt.equals(status.readersBuilt);
+            return readersBuilt.equals(status.readersBuilt) && valuesLoaded == status.valuesLoaded;
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(super.hashCode(), readersBuilt);
+            return Objects.hash(super.hashCode(), readersBuilt, valuesLoaded);
         }
 
         @Override
@@ -710,6 +735,4 @@ public BlockLoader.AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int
             return factory.newAggregateMetricDoubleBlockBuilder(count);
         }
     }
-
-    // TODO tests that mix source loaded fields and doc values in the same block
 }