Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions docs/changelog/123396.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
pr: 123396
summary: Add initial grammar and planning for RRF (snapshot)
area: ES|QL
type: feature
issues: []
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
/*
* 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;

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));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: this is currently configurable in _search, so we probably need to expose it as an option in the future here too

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right, we need to make the rank constant configurable.
This is added as a separate feature in the meta issue #123391
It will require a syntax change for RRF, so I'd like to keep it separate for now.

}

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();
}
});
}
}
40 changes: 40 additions & 0 deletions x-pack/plugin/esql/qa/testFixtures/src/main/resources/rrf.csv-spec
Original file line number Diff line number Diff line change
@@ -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 )
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we can have some queries with disjunctions in the where clause of each fork leg that will be great, just to add a bit more complexity to make sure it works as expected. There are some queries with disjunctions in the match function and operator's csvtests, that can be used as a reference.

( 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
;
Original file line number Diff line number Diff line change
Expand Up @@ -350,6 +350,32 @@ 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" )
( WHERE content:"dog" )
| RRF
| 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.032266458495966696, List.of("fork1", "fork2")),
List.of(3, "This dog is really brown", 0.01639344262295082, "fork2"),
List.of(4, "The dog is brown but this document is very very long", 0.016129032258064516, "fork2")
);
assertValues(resp.values(), expectedValues);
}
}

public void testThreeSubQueries() {
var query = """
FROM test
Expand Down
1 change: 1 addition & 0 deletions x-pack/plugin/esql/src/main/antlr/EsqlBaseLexer.g4
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ import ChangePoint,
Metrics,
MvExpand,
Project,
Rrf,
Rename,
Show,
UnknownCommand;
Expand Down
Loading