Merge branch 'adj-ivf-postings-list-building' of github.com:benwtrent/elasticsearch into adj-ivf-postings-list-building

Author: benwtrent

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

@@ -275,7 +275,7 @@ For the complete Elasticsearch documentation visit
 https://www.elastic.co/guide/en/elasticsearch/reference/current/index.html[elastic.co].
 
 For information about our documentation processes, see the
-xref:docs/README.asciidoc[docs README].
+xref:https://github.com/elastic/elasticsearch/blob/main/docs/README.md[docs README].
 
 [[examples]]
 == Examples and guides

README.asciidoc

@@ -70,11 +70,11 @@ public class QueryPlanningBenchmark {
     private EsqlParser defaultParser;
     private Analyzer manyFieldsAnalyzer;
     private LogicalPlanOptimizer defaultOptimizer;
+    private Configuration config;
 
     @Setup
     public void setup() {
-
-        var config = new Configuration(
+        this.config = new Configuration(
             DateUtils.UTC,
             Locale.US,
             null,
@@ -116,7 +116,7 @@ public void setup() {
     }
 
     private LogicalPlan plan(EsqlParser parser, Analyzer analyzer, LogicalPlanOptimizer optimizer, String query) {
-        var parsed = parser.createStatement(query, new QueryParams(), telemetry);
+        var parsed = parser.createStatement(query, new QueryParams(), telemetry, config);
         var analyzed = analyzer.analyze(parsed);
         var optimized = optimizer.optimize(analyzed);
         return optimized;

benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/QueryPlanningBenchmark.java

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-package org.elasticsearch.benchmark.compute.operator;
+package org.elasticsearch.benchmark._nightly.esql;
 
 import org.apache.lucene.document.FieldType;
 import org.apache.lucene.document.NumericDocValuesField;
@@ -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;
@@ -84,10 +85,19 @@
 @State(Scope.Thread)
 @Fork(1)
 public class ValuesSourceReaderBenchmark {
+    private static final String[] SUPPORTED_LAYOUTS = new String[] { "in_order", "shuffled", "shuffled_singles" };
+    private static final String[] SUPPORTED_NAMES = new String[] {
+        "long",
+        "int",
+        "double",
+        "keyword",
+        "stored_keyword",
+        "3_stored_keywords",
+        "keyword_mv" };
+
     private static final int BLOCK_LENGTH = 16 * 1024;
     private static final int INDEX_SIZE = 10 * BLOCK_LENGTH;
     private static final int COMMIT_INTERVAL = 500;
-    private static final BigArrays BIG_ARRAYS = BigArrays.NON_RECYCLING_INSTANCE;
     private static final BlockFactory blockFactory = BlockFactory.getInstance(
         new NoopCircuitBreaker("noop"),
         BigArrays.NON_RECYCLING_INSTANCE
@@ -103,8 +113,8 @@ static void selfTest() {
             ValuesSourceReaderBenchmark benchmark = new ValuesSourceReaderBenchmark();
             benchmark.setupIndex();
             try {
-                for (String layout : ValuesSourceReaderBenchmark.class.getField("layout").getAnnotationsByType(Param.class)[0].value()) {
-                    for (String name : ValuesSourceReaderBenchmark.class.getField("name").getAnnotationsByType(Param.class)[0].value()) {
+                for (String layout : ValuesSourceReaderBenchmark.SUPPORTED_LAYOUTS) {
+                    for (String name : ValuesSourceReaderBenchmark.SUPPORTED_NAMES) {
                         benchmark.layout = layout;
                         benchmark.name = name;
                         try {
@@ -118,7 +128,7 @@ static void selfTest() {
             } finally {
                 benchmark.teardownIndex();
             }
-        } catch (IOException | NoSuchFieldException e) {
+        } catch (IOException e) {
             throw new AssertionError(e);
         }
     }
@@ -320,10 +330,10 @@ public FieldNamesFieldMapper.FieldNamesFieldType fieldNames() {
      *     each page has a single document rather than {@code BLOCK_SIZE} docs.</li>
      * </ul>
      */
-    @Param({ "in_order", "shuffled", "shuffled_singles" })
+    @Param({ "in_order", "shuffled" })
     public String layout;
 
-    @Param({ "long", "int", "double", "keyword", "stored_keyword", "3_stored_keywords" })
+    @Param({ "long", "keyword", "stored_keyword", "keyword_mv" })
     public String name;
 
     private Directory directory;
@@ -343,7 +353,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();
@@ -389,6 +399,22 @@ public void benchmark() {
                         }
                     }
                 }
+                case "keyword_mv" -> {
+                    BytesRef scratch = new BytesRef();
+                    BytesRefBlock values = op.getOutput().<BytesRefBlock>getBlock(1);
+                    for (int p = 0; p < values.getPositionCount(); p++) {
+                        int count = values.getValueCount(p);
+                        if (count > 0) {
+                            int first = values.getFirstValueIndex(p);
+                            for (int i = 0; i < count; i++) {
+                                BytesRef r = values.getBytesRef(first + i, scratch);
+                                r.offset++;
+                                r.length--;
+                                sum += Integer.parseInt(r.utf8ToString());
+                            }
+                        }
+                    }
+                }
             }
         }
         long expected = 0;
@@ -398,6 +424,16 @@ public void benchmark() {
                     expected += i % 1000;
                 }
                 break;
+            case "keyword_mv":
+                for (int i = 0; i < INDEX_SIZE; i++) {
+                    int v1 = i % 1000;
+                    expected += v1;
+                    int v2 = i % 500;
+                    if (v1 != v2) {
+                        expected += v2;
+                    }
+                }
+                break;
             case "3_stored_keywords":
                 for (int i = 0; i < INDEX_SIZE; i++) {
                     expected += 3 * (i % 1000);
@@ -411,7 +447,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;
@@ -452,7 +488,9 @@ private void setupIndex() throws IOException {
                         new StoredField("double", (double) i),
                         new KeywordFieldMapper.KeywordField("keyword_1", new BytesRef(c + i % 1000), keywordFieldType),
                         new KeywordFieldMapper.KeywordField("keyword_2", new BytesRef(c + i % 1000), keywordFieldType),
-                        new KeywordFieldMapper.KeywordField("keyword_3", new BytesRef(c + i % 1000), keywordFieldType)
+                        new KeywordFieldMapper.KeywordField("keyword_3", new BytesRef(c + i % 1000), keywordFieldType),
+                        new KeywordFieldMapper.KeywordField("keyword_mv", new BytesRef(c + i % 1000), keywordFieldType),
+                        new KeywordFieldMapper.KeywordField("keyword_mv", new BytesRef(c + i % 500), keywordFieldType)
                     )
                 );
                 if (i % COMMIT_INTERVAL == 0) {

…perator/ValuesSourceReaderBenchmark.java …ly/esql/ValuesSourceReaderBenchmark.javabenchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesSourceReaderBenchmark.java renamed to benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/ValuesSourceReaderBenchmark.java benchmarks/src/main/java/org/elasticsearch/benchmark/compute/operator/ValuesSourceReaderBenchmark.java renamed to benchmarks/src/main/java/org/elasticsearch/benchmark/_nightly/esql/ValuesSourceReaderBenchmark.java

@@ -125,7 +125,12 @@ 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() {
+                return super.getOutput();
+            }
+        };
     }
 
     private static AggregatorFunctionSupplier supplier(String dataType) {

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

@@ -8,12 +8,14 @@
  */
 package org.elasticsearch.benchmark.vector;
 
+import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.lucene.store.MMapDirectory;
 import org.apache.lucene.util.VectorUtil;
+import org.apache.lucene.util.quantization.OptimizedScalarQuantizer;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.simdvec.ES91Int4VectorsScorer;
@@ -52,20 +54,26 @@ public class Int4ScorerBenchmark {
         LogConfigurator.configureESLogging(); // native access requires logging to be initialized
     }
 
-    @Param({ "384", "702", "1024" })
+    @Param({ "384", "782", "1024" })
     int dims;
 
-    int numVectors = 200;
-    int numQueries = 10;
+    int numVectors = 20 * ES91Int4VectorsScorer.BULK_SIZE;
+    int numQueries = 5;
 
     byte[] scratch;
     byte[][] binaryVectors;
     byte[][] binaryQueries;
+    float[] scores = new float[ES91Int4VectorsScorer.BULK_SIZE];
+
+    float[] scratchFloats = new float[3];
 
     ES91Int4VectorsScorer scorer;
     Directory dir;
     IndexInput in;
 
+    OptimizedScalarQuantizer.QuantizationResult queryCorrections;
+    float centroidDp;
+
     @Setup
     public void setup() throws IOException {
         binaryVectors = new byte[numVectors][dims];
@@ -77,9 +85,19 @@ public void setup() throws IOException {
                     binaryVector[i] = (byte) ThreadLocalRandom.current().nextInt(16);
                 }
                 out.writeBytes(binaryVector, 0, binaryVector.length);
+                ThreadLocalRandom.current().nextBytes(binaryVector);
+                out.writeBytes(binaryVector, 0, 14); // corrections
             }
         }
 
+        queryCorrections = new OptimizedScalarQuantizer.QuantizationResult(
+            ThreadLocalRandom.current().nextFloat(),
+            ThreadLocalRandom.current().nextFloat(),
+            ThreadLocalRandom.current().nextFloat(),
+            Short.toUnsignedInt((short) ThreadLocalRandom.current().nextInt())
+        );
+        centroidDp = ThreadLocalRandom.current().nextFloat();
+
         in = dir.openInput("vectors", IOContext.DEFAULT);
         binaryQueries = new byte[numVectors][dims];
         for (byte[] binaryVector : binaryVectors) {
@@ -105,18 +123,66 @@ public void scoreFromArray(Blackhole bh) throws IOException {
             in.seek(0);
             for (int i = 0; i < numVectors; i++) {
                 in.readBytes(scratch, 0, dims);
-                bh.consume(VectorUtil.int4DotProduct(binaryQueries[j], scratch));
+                int dp = VectorUtil.int4DotProduct(binaryQueries[j], scratch);
+                in.readFloats(scratchFloats, 0, 3);
+                float score = scorer.applyCorrections(
+                    queryCorrections.lowerInterval(),
+                    queryCorrections.upperInterval(),
+                    queryCorrections.quantizedComponentSum(),
+                    queryCorrections.additionalCorrection(),
+                    VectorSimilarityFunction.EUCLIDEAN,
+                    centroidDp, // assuming no centroid dot product for this benchmark
+                    scratchFloats[0],
+                    scratchFloats[1],
+                    Short.toUnsignedInt(in.readShort()),
+                    scratchFloats[2],
+                    dp
+                );
+                bh.consume(score);
             }
         }
     }
 
     @Benchmark
     @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
-    public void scoreFromMemorySegmentOnlyVector(Blackhole bh) throws IOException {
+    public void scoreFromMemorySegment(Blackhole bh) throws IOException {
         for (int j = 0; j < numQueries; j++) {
             in.seek(0);
             for (int i = 0; i < numVectors; i++) {
-                bh.consume(scorer.int4DotProduct(binaryQueries[j]));
+                bh.consume(
+                    scorer.score(
+                        binaryQueries[j],
+                        queryCorrections.lowerInterval(),
+                        queryCorrections.upperInterval(),
+                        queryCorrections.quantizedComponentSum(),
+                        queryCorrections.additionalCorrection(),
+                        VectorSimilarityFunction.EUCLIDEAN,
+                        centroidDp
+                    )
+                );
+            }
+        }
+    }
+
+    @Benchmark
+    @Fork(jvmArgsPrepend = { "--add-modules=jdk.incubator.vector" })
+    public void scoreFromMemorySegmentBulk(Blackhole bh) throws IOException {
+        for (int j = 0; j < numQueries; j++) {
+            in.seek(0);
+            for (int i = 0; i < numVectors; i += ES91Int4VectorsScorer.BULK_SIZE) {
+                scorer.scoreBulk(
+                    binaryQueries[j],
+                    queryCorrections.lowerInterval(),
+                    queryCorrections.upperInterval(),
+                    queryCorrections.quantizedComponentSum(),
+                    queryCorrections.additionalCorrection(),
+                    VectorSimilarityFunction.EUCLIDEAN,
+                    centroidDp,
+                    scores
+                );
+                for (float score : scores) {
+                    bh.consume(score);
+                }
             }
         }
     }