ES|QL: Add initial grammar and planning for RRF (snapshot)

docs/changelog/123396.yaml

@@ -0,0 +1,5 @@
+pr: 123396
+summary: Add initial grammar and planning for RRF (snapshot)
+area: ES|QL
+type: feature
+issues: []

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/RrfScoreEvalOperator.java

@@ -0,0 +1,113 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.compute.operator;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.DoubleVector;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Releasables;
+
+import java.util.HashMap;
+
+/**
+ * Updates the score column with new scores using the RRF formula.
+ * Receives the position of the score and fork columns.
+ * The new score we assign to each row is equal to {@code 1 / (rank_constant + row_number)}.
+ * We use the fork discriminator column to determine the {@code row_number} for each row.
+ */
+public class RrfScoreEvalOperator implements Operator {
+
+    public record Factory(int forkPosition, int scorePosition) implements OperatorFactory {
+        @Override
+        public Operator get(DriverContext driverContext) {
+            return new RrfScoreEvalOperator(forkPosition, scorePosition);
+        }
+
+        @Override
+        public String describe() {
+            return "RrfScoreEvalOperator";
+        }
+
+    }
+
+    private final int scorePosition;
+    private final int forkPosition;
+
+    private boolean finished = false;
+    private Page prev = null;
+
+    private HashMap<String, Integer> counters = new HashMap<>();
+
+    public RrfScoreEvalOperator(int forkPosition, int scorePosition) {
+        this.scorePosition = scorePosition;
+        this.forkPosition = forkPosition;
+    }
+
+    @Override
+    public boolean needsInput() {
+        return prev == null && finished == false;
+    }
+
+    @Override
+    public void addInput(Page page) {
+        assert prev == null : "has pending input page";
+        prev = page;
+    }
+
+    @Override
+    public void finish() {
+        finished = true;
+    }
+
+    @Override
+    public boolean isFinished() {
+        return finished && prev == null;
+    }
+
+    @Override
+    public Page getOutput() {
+        Page page = prev;
+
+        BytesRefBlock forkBlock = (BytesRefBlock) page.getBlock(forkPosition);
+
+        DoubleVector.Builder scores = forkBlock.blockFactory().newDoubleVectorBuilder(forkBlock.getPositionCount());
+
+        for (int i = 0; i < page.getPositionCount(); i++) {
+            String fork = forkBlock.getBytesRef(i, new BytesRef()).utf8ToString();
+
+            int rank = counters.getOrDefault(fork, 1);
+            counters.put(fork, rank + 1);
+            scores.appendDouble(1.0 / (60 + rank));
+        }
+
+        Block scoreBlock = scores.build().asBlock();
+        page = page.appendBlock(scoreBlock);
+
+        int[] projections = new int[page.getBlockCount() - 1];
+
+        for (int i = 0; i < page.getBlockCount() - 1; i++) {
+            projections[i] = i == scorePosition ? page.getBlockCount() - 1 : i;
+        }
+
+        page = page.projectBlocks(projections);
+
+        prev = null;
+        return page;
+    }
+
+    @Override
+    public void close() {
+        Releasables.closeExpectNoException(() -> {
+            if (prev != null) {
+                prev.releaseBlocks();
+            }
+        });
+    }
+}

x-pack/plugin/esql/qa/testFixtures/src/main/resources/rrf.csv-spec

@@ -0,0 +1,40 @@
+//
+// CSV spec for RRF command
+//
+
+simpleRrf
+required_capability: fork
+required_capability: rrf
+required_capability: match_operator_colon
+
+FROM employees METADATA _id, _index, _score
+| FORK ( WHERE emp_no:10001 )
+       ( WHERE emp_no:10002 )
+| RRF
+| KEEP _score, _fork, emp_no
+;
+
+_score:double       | _fork:keyword  | emp_no:integer
+0.01639344262295082 | fork1          | 10001
+0.01639344262295082 | fork2          | 10002
+;
+
+rrfWithMatchAndScore
+required_capability: fork
+required_capability: rrf
+required_capability: match_operator_colon
+
+FROM books METADATA _id, _index, _score
+| FORK ( WHERE title:"Tolkien" | SORT _score DESC | LIMIT 3 )
+       ( WHERE author:"Tolkien" | SORT _score DESC | LIMIT 3 )
+| RRF
+| KEEP _score, _fork, _id
+;
+
+_score:double        | _fork:keyword  | _id:keyword
+0.03225806451612903  | [fork1, fork2] | 26
+0.01639344262295082  | fork2          | 18
+0.01639344262295082  | fork1          | 36
+0.015873015873015872 | fork1          | 56
+0.015873015873015872 | fork2          | 59
+;

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ForkIT.java

@@ -350,6 +350,33 @@ public void testScoringKeepAndSort() {
         }
     }
 
+    public void testRrf() {
+        assumeTrue("requires RRF capability", EsqlCapabilities.Cap.RRF.isEnabled());
+
+        var query = """
+            FROM test METADATA _score, _id, _index
+            | WHERE id > 2
+            | FORK
+               ( WHERE content:"fox" | SORT _score, _id DESC )
+               ( WHERE content:"dog" | SORT _score, _id DESC )
+            | RRF
+            | EVAL _score = round(_score, 4)
+            | KEEP id, content, _score, _fork
+            """;
+        try (var resp = run(query)) {
+            System.out.println("response=" + resp);
+            assertColumnNames(resp.columns(), List.of("id", "content", "_score", "_fork"));
+            assertColumnTypes(resp.columns(), List.of("integer", "keyword", "double", "keyword"));
+            assertThat(getValuesList(resp.values()).size(), equalTo(3));
+            Iterable<Iterable<Object>> expectedValues = List.of(
+                List.of(6, "The quick brown fox jumps over the lazy dog", 0.0325, List.of("fork1", "fork2")),
+                List.of(4, "The dog is brown but this document is very very long", 0.0164, "fork2"),
+                List.of(3, "This dog is really brown", 0.0159, "fork2")
+            );
+            assertValues(resp.values(), expectedValues);
+        }
+    }
+
     public void testThreeSubQueries() {
         var query = """
             FROM test

x-pack/plugin/esql/src/main/antlr/EsqlBaseLexer.g4

@@ -65,6 +65,7 @@ import ChangePoint,
        Metrics,
        MvExpand,
        Project,
+       Rrf,
        Rename,
        Show,
        UnknownCommand;