Merge branch 'main' into ms-graph-docs

Author: richard-dennehy

.buildkite/pipelines/periodic.template.yml

@@ -204,6 +204,17 @@ steps:
           image: family/elasticsearch-ubuntu-2404
           machineType: n2-standard-8
           buildDirectory: /dev/shm/bk
+      - label: third-party / ms-graph
+        command: |
+          .ci/scripts/run-gradle.sh msGraphThirdPartyTest
+        env:
+          USE_3RD_PARTY_MS_GRAPH_CREDENTIALS: "true"
+        timeout_in_minutes: 30
+        agents:
+          provider: gcp
+          image: family/elasticsearch-ubuntu-2404
+          machineType: n2-standard-8
+          buildDirectory: /dev/shm/bk
   - group: lucene-compat
     steps:
       - label: "{{matrix.LUCENE_VERSION}} / lucene-compat"

.buildkite/pipelines/periodic.yml

@@ -642,6 +642,17 @@ steps:
           image: family/elasticsearch-ubuntu-2404
           machineType: n2-standard-8
           buildDirectory: /dev/shm/bk
+      - label: third-party / ms-graph
+        command: |
+          .ci/scripts/run-gradle.sh msGraphThirdPartyTest
+        env:
+          USE_3RD_PARTY_MS_GRAPH_CREDENTIALS: "true"
+        timeout_in_minutes: 30
+        agents:
+          provider: gcp
+          image: family/elasticsearch-ubuntu-2404
+          machineType: n2-standard-8
+          buildDirectory: /dev/shm/bk
   - group: lucene-compat
     steps:
       - label: "{{matrix.LUCENE_VERSION}} / lucene-compat"

.buildkite/scripts/third-party-test-credentials.sh

@@ -48,6 +48,23 @@ if [[ "${USE_3RD_PARTY_GCS_CREDENTIALS:-}" == "true" ]]; then
   .buildkite/scripts/third-party-test-credentials.gcs.sh "$google_storage_service_account"
 fi
 
+if [[ "${USE_3RD_PARTY_MS_GRAPH_CREDENTIALS:-}" == "true" ]]; then
+  json=$(vault read -format=json secret/ci/elastic-elasticsearch/ms_graph_thirdparty_test_creds)
 
+  MS_GRAPH_TENANT_ID=$(echo "$json" | jq -r .data.tenant_id)
+  export ms_graph_tenant_id="$MS_GRAPH_TENANT_ID"
+
+  MS_GRAPH_CLIENT_ID=$(echo "$json" | jq -r .data.client_id)
+  export ms_graph_client_id="$MS_GRAPH_CLIENT_ID"
+
+  MS_GRAPH_CLIENT_SECRET=$(echo "$json" | jq -r .data.client_secret)
+  export ms_graph_client_secret="$MS_GRAPH_CLIENT_SECRET"
+
+  MS_GRAPH_USERNAME=$(echo "$json" | jq -r .data.username)
+  export ms_graph_username="$MS_GRAPH_USERNAME"
+
+  MS_GRAPH_GROUP_ID=$(echo "$json" | jq -r .data.group_id)
+  export ms_graph_group_id="$MS_GRAPH_GROUP_ID"
+fi
 
 unset json

.github/workflows/docs-preview-comment.yml

@@ -95,7 +95,8 @@ static void selfTest() {
         try {
             for (String groups : ValuesAggregatorBenchmark.class.getField("groups").getAnnotationsByType(Param.class)[0].value()) {
                 for (String dataType : ValuesAggregatorBenchmark.class.getField("dataType").getAnnotationsByType(Param.class)[0].value()) {
-                    run(Integer.parseInt(groups), dataType, 10);
+                    run(Integer.parseInt(groups), dataType, 10, 0);
+                    run(Integer.parseInt(groups), dataType, 10, 1);
                 }
             }
         } catch (NoSuchFieldException e) {
@@ -113,7 +114,10 @@ static void selfTest() {
     @Param({ BYTES_REF, INT, LONG })
     public String dataType;
 
-    private static Operator operator(DriverContext driverContext, int groups, String dataType) {
+    @Param({ "0", "1" })
+    public int numOrdinalMerges;
+
+    private static Operator operator(DriverContext driverContext, int groups, String dataType, int numOrdinalMerges) {
         if (groups == 1) {
             return new AggregationOperator(
                 List.of(supplier(dataType).aggregatorFactory(AggregatorMode.SINGLE, List.of(0)).apply(driverContext)),
@@ -125,7 +129,24 @@ private static Operator operator(DriverContext driverContext, int groups, String
             List.of(supplier(dataType).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(1))),
             () -> BlockHash.build(groupSpec, driverContext.blockFactory(), 16 * 1024, false),
             driverContext
-        );
+        ) {
+            @Override
+            public Page getOutput() {
+                mergeOrdinal();
+                return super.getOutput();
+            }
+
+            // simulate OrdinalsGroupingOperator
+            void mergeOrdinal() {
+                var merged = supplier(dataType).groupingAggregatorFactory(AggregatorMode.SINGLE, List.of(1)).apply(driverContext);
+                for (int i = 0; i < numOrdinalMerges; i++) {
+                    for (int p = 0; p < groups; p++) {
+                        merged.addIntermediateRow(p, aggregators.getFirst(), p);
+                    }
+                }
+                aggregators.set(0, merged);
+            }
+        };
     }
 
     private static AggregatorFunctionSupplier supplier(String dataType) {
@@ -331,12 +352,12 @@ private static Block groupingBlock(int groups) {
 
     @Benchmark
     public void run() {
-        run(groups, dataType, OP_COUNT);
+        run(groups, dataType, OP_COUNT, numOrdinalMerges);
     }
 
-    private static void run(int groups, String dataType, int opCount) {
+    private static void run(int groups, String dataType, int opCount, int numOrdinalMerges) {
         DriverContext driverContext = driverContext();
-        try (Operator operator = operator(driverContext, groups, dataType)) {
+        try (Operator operator = operator(driverContext, groups, dataType, numOrdinalMerges)) {
             Page page = page(groups, dataType);
             for (int i = 0; i < opCount; i++) {
                 operator.addInput(page.shallowCopy());

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesAggregatorBenchmark.java

@@ -41,7 +41,8 @@
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.lucene.LuceneSourceOperator;
 import org.elasticsearch.compute.lucene.ShardRefCounted;
-import org.elasticsearch.compute.lucene.ValuesSourceReaderOperator;
+import org.elasticsearch.compute.lucene.read.ValuesSourceReaderOperator;
+import org.elasticsearch.compute.lucene.read.ValuesSourceReaderOperatorStatus;
 import org.elasticsearch.compute.operator.topn.TopNOperator;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.index.IndexSettings;
@@ -343,7 +344,7 @@ public void benchmark() {
         );
         long sum = 0;
         for (Page page : pages) {
-            op.addInput(page);
+            op.addInput(page.shallowCopy());
             switch (name) {
                 case "long" -> {
                     LongVector values = op.getOutput().<LongBlock>getBlock(1).asVector();
@@ -411,7 +412,7 @@ public void benchmark() {
             throw new AssertionError("[" + layout + "][" + name + "] expected [" + expected + "] but was [" + sum + "]");
         }
         boolean foundStoredFieldLoader = false;
-        ValuesSourceReaderOperator.Status status = (ValuesSourceReaderOperator.Status) op.status();
+        ValuesSourceReaderOperatorStatus status = (ValuesSourceReaderOperatorStatus) op.status();
         for (Map.Entry<String, Integer> e : status.readersBuilt().entrySet()) {
             if (e.getKey().indexOf("stored_fields") >= 0) {
                 foundStoredFieldLoader = true;

benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesSourceReaderBenchmark.java

@@ -24,6 +24,8 @@
 import org.apache.lucene.util.quantization.ScalarQuantizer;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 import org.elasticsearch.simdvec.VectorScorerFactory;
 import org.openjdk.jmh.annotations.Benchmark;
 import org.openjdk.jmh.annotations.BenchmarkMode;
@@ -55,24 +57,27 @@
 /**
  * Benchmark that compares various scalar quantized vector similarity function
  * implementations;: scalar, lucene's panama-ized, and Elasticsearch's native.
- * Run with ./gradlew -p benchmarks run --args 'VectorScorerBenchmark'
+ * Run with ./gradlew -p benchmarks run --args 'Int7uScorerBenchmark'
  */
-public class VectorScorerBenchmark {
+public class Int7uScorerBenchmark {
 
     static {
         LogConfigurator.configureESLogging(); // native access requires logging to be initialized
+        if (supportsHeapSegments() == false) {
+            final Logger LOG = LogManager.getLogger(Int7uScorerBenchmark.class);
+            LOG.warn("*Query targets cannot run on " + "JDK " + Runtime.version());
+        }
     }
 
     @Param({ "96", "768", "1024" })
-    int dims;
-    int size = 2; // there are only two vectors to compare
+    public int dims;
+    final int size = 2; // there are only two vectors to compare
 
     Directory dir;
     IndexInput in;
     VectorScorerFactory factory;
 
-    byte[] vec1;
-    byte[] vec2;
+    byte[] vec1, vec2;
     float vec1Offset;
     float vec2Offset;
     float scoreCorrectionConstant;
@@ -130,47 +135,16 @@ public void setup() throws IOException {
         nativeSqrScorer = factory.getInt7SQVectorScorerSupplier(EUCLIDEAN, in, values, scoreCorrectionConstant).get().scorer();
         nativeSqrScorer.setScoringOrdinal(0);
 
-        // setup for getInt7SQVectorScorer / query vector scoring
-        float[] queryVec = new float[dims];
-        for (int i = 0; i < dims; i++) {
-            queryVec[i] = ThreadLocalRandom.current().nextFloat();
-        }
-        luceneDotScorerQuery = luceneScorer(values, VectorSimilarityFunction.DOT_PRODUCT, queryVec);
-        nativeDotScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.DOT_PRODUCT, values, queryVec).get();
-        luceneSqrScorerQuery = luceneScorer(values, VectorSimilarityFunction.EUCLIDEAN, queryVec);
-        nativeSqrScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.EUCLIDEAN, values, queryVec).get();
-
-        // sanity
-        var f1 = dotProductLucene();
-        var f2 = dotProductNative();
-        var f3 = dotProductScalar();
-        if (f1 != f2) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
-        }
-        if (f1 != f3) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
-        }
-        // square distance
-        f1 = squareDistanceLucene();
-        f2 = squareDistanceNative();
-        f3 = squareDistanceScalar();
-        if (f1 != f2) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "native[" + f2 + "]");
-        }
-        if (f1 != f3) {
-            throw new AssertionError("lucene[" + f1 + "] != " + "scalar[" + f3 + "]");
-        }
-
-        var q1 = dotProductLuceneQuery();
-        var q2 = dotProductNativeQuery();
-        if (q1 != q2) {
-            throw new AssertionError("query: lucene[" + q1 + "] != " + "native[" + q2 + "]");
-        }
-
-        var sqr1 = squareDistanceLuceneQuery();
-        var sqr2 = squareDistanceNativeQuery();
-        if (sqr1 != sqr2) {
-            throw new AssertionError("query: lucene[" + q1 + "] != " + "native[" + q2 + "]");
+        if (supportsHeapSegments()) {
+            // setup for getInt7SQVectorScorer / query vector scoring
+            float[] queryVec = new float[dims];
+            for (int i = 0; i < dims; i++) {
+                queryVec[i] = ThreadLocalRandom.current().nextFloat();
+            }
+            luceneDotScorerQuery = luceneScorer(values, VectorSimilarityFunction.DOT_PRODUCT, queryVec);
+            nativeDotScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.DOT_PRODUCT, values, queryVec).get();
+            luceneSqrScorerQuery = luceneScorer(values, VectorSimilarityFunction.EUCLIDEAN, queryVec);
+            nativeSqrScorerQuery = factory.getInt7SQVectorScorer(VectorSimilarityFunction.EUCLIDEAN, values, queryVec).get();
         }
     }
 
@@ -242,6 +216,10 @@ public float squareDistanceNativeQuery() throws IOException {
         return nativeSqrScorerQuery.score(1);
     }
 
+    static boolean supportsHeapSegments() {
+        return Runtime.version().feature() >= 22;
+    }
+
     QuantizedByteVectorValues vectorValues(int dims, int size, IndexInput in, VectorSimilarityFunction sim) throws IOException {
         var sq = new ScalarQuantizer(0.1f, 0.9f, (byte) 7);
         var slice = in.slice("values", 0, in.length());