Skip to content

Commit b926d41

Browse files
nik9000nik9000
authored
ESQL: Support runtime doubles (ESQL-606)
This builds on the work I did in ESQL-591 to add support for runtime `double` fields. And it adds simple tests for runtime `keyword` fields as well. Those already worked without any change.
1 parent fb17603 commit b926d41

File tree

2 files changed

+138
-26
lines changed

2 files changed

+138
-26
lines changed

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

Lines changed: 46 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,10 @@ public static BlockDocValuesReader createBlockReader(
6363
ValuesSource.Numeric numericVS = (ValuesSource.Numeric) valuesSource;
6464
if (numericVS.isFloatingPoint()) {
6565
final SortedNumericDoubleValues doubleValues = numericVS.doubleValues(leafReaderContext);
66+
final NumericDoubleValues singleton = FieldData.unwrapSingleton(doubleValues);
67+
if (singleton != null) {
68+
return new DoubleSingletonValuesReader(singleton);
69+
}
6670
return new DoubleValuesReader(doubleValues);
6771
} else {
6872
final SortedNumericDocValues longValues = numericVS.longValues(leafReaderContext);
@@ -152,12 +156,12 @@ public int docID() {
152156
}
153157
}
154158

155-
private static class DoubleValuesReader extends BlockDocValuesReader {
159+
private static class DoubleSingletonValuesReader extends BlockDocValuesReader {
156160
private final NumericDoubleValues numericDocValues;
157161
private int docID = -1;
158162

159-
DoubleValuesReader(SortedNumericDoubleValues numericDocValues) {
160-
this.numericDocValues = FieldData.unwrapSingleton(numericDocValues);
163+
DoubleSingletonValuesReader(NumericDoubleValues numericDocValues) {
164+
this.numericDocValues = numericDocValues;
161165
}
162166

163167
@Override
@@ -188,6 +192,45 @@ public int docID() {
188192
}
189193
}
190194

195+
private static class DoubleValuesReader extends BlockDocValuesReader {
196+
private final SortedNumericDoubleValues numericDocValues;
197+
private int docID = -1;
198+
199+
DoubleValuesReader(SortedNumericDoubleValues numericDocValues) {
200+
this.numericDocValues = numericDocValues;
201+
}
202+
203+
@Override
204+
public Block readValues(IntVector docs) throws IOException {
205+
final int positionCount = docs.getPositionCount();
206+
var blockBuilder = DoubleBlock.newBlockBuilder(positionCount);
207+
int lastDoc = -1;
208+
for (int i = 0; i < positionCount; i++) {
209+
int doc = docs.getInt(i);
210+
// docs within same block must be in order
211+
if (lastDoc >= doc) {
212+
throw new IllegalStateException("docs within same block must be in order");
213+
}
214+
if (numericDocValues.advanceExact(doc)) {
215+
if (numericDocValues.docValueCount() != 1) {
216+
throw new UnsupportedOperationException("only single valued fields supported for now");
217+
}
218+
blockBuilder.appendDouble(numericDocValues.nextValue());
219+
} else {
220+
blockBuilder.appendNull();
221+
}
222+
lastDoc = doc;
223+
this.docID = doc;
224+
}
225+
return blockBuilder.build();
226+
}
227+
228+
@Override
229+
public int docID() {
230+
return docID;
231+
}
232+
}
233+
191234
private static class BytesValuesReader extends BlockDocValuesReader {
192235
private int docID = -1;
193236
private final SortedBinaryDocValues binaryDV;

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

Lines changed: 92 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,11 @@
1212
import org.elasticsearch.index.mapper.OnScriptError;
1313
import org.elasticsearch.plugins.Plugin;
1414
import org.elasticsearch.plugins.ScriptPlugin;
15+
import org.elasticsearch.script.DoubleFieldScript;
1516
import org.elasticsearch.script.LongFieldScript;
1617
import org.elasticsearch.script.ScriptContext;
1718
import org.elasticsearch.script.ScriptEngine;
19+
import org.elasticsearch.script.StringFieldScript;
1820
import org.elasticsearch.search.lookup.SearchLookup;
1921
import org.elasticsearch.test.ESIntegTestCase;
2022
import org.elasticsearch.xcontent.XContentBuilder;
@@ -37,18 +39,48 @@
3739
*/
3840
@ESIntegTestCase.ClusterScope(scope = SUITE, numDataNodes = 1, numClientNodes = 0, supportsDedicatedMasters = false) // ESQL is single node
3941
public class EsqlActionRuntimeFieldIT extends ESIntegTestCase {
42+
private static final int SIZE = 5000;
43+
4044
@Override
4145
protected Collection<Class<? extends Plugin>> nodePlugins() {
42-
return List.of(EsqlPlugin.class, PausableFieldPlugin.class);
46+
return List.of(EsqlPlugin.class, TestRuntimeFieldPlugin.class);
47+
}
48+
49+
public void testLong() throws InterruptedException, IOException {
50+
createIndexWithConstRuntimeField("long");
51+
EsqlQueryResponse response = EsqlActionIT.run("from test | stats sum(const)", Settings.EMPTY);
52+
assertThat(response.values(), equalTo(List.of(List.of((long) SIZE))));
53+
}
54+
55+
public void testDouble() throws InterruptedException, IOException {
56+
createIndexWithConstRuntimeField("double");
57+
EsqlQueryResponse response = EsqlActionIT.run("from test | stats sum(const)", Settings.EMPTY);
58+
assertThat(response.values(), equalTo(List.of(List.of((double) SIZE))));
59+
}
60+
61+
public void testKeyword() throws InterruptedException, IOException {
62+
createIndexWithConstRuntimeField("keyword");
63+
EsqlQueryResponse response = EsqlActionIT.run("from test | project const | limit 1", Settings.EMPTY);
64+
assertThat(response.values(), equalTo(List.of(List.of("const"))));
65+
}
66+
67+
/**
68+
* Test grouping by runtime keyword which requires disabling the ordinals
69+
* optimization available to more keyword fields.
70+
*/
71+
public void testKeywordBy() throws InterruptedException, IOException {
72+
createIndexWithConstRuntimeField("keyword");
73+
EsqlQueryResponse response = EsqlActionIT.run("from test | stats max(foo) by const", Settings.EMPTY);
74+
assertThat(response.values(), equalTo(List.of(List.of(SIZE - 1L, "const"))));
4375
}
4476

45-
public void testTask() throws InterruptedException, IOException {
77+
private void createIndexWithConstRuntimeField(String type) throws InterruptedException, IOException {
4678
XContentBuilder mapping = JsonXContent.contentBuilder().startObject();
4779
mapping.startObject("runtime");
4880
{
49-
mapping.startObject("pause_me");
81+
mapping.startObject("const");
5082
{
51-
mapping.field("type", "long");
83+
mapping.field("type", type);
5284
mapping.startObject("script").field("source", "").field("lang", "dummy").endObject();
5385
}
5486
mapping.endObject();
@@ -61,11 +93,9 @@ public void testTask() throws InterruptedException, IOException {
6193
indexRequests.add(client().prepareIndex("test").setId(Integer.toString(i)).setSource("foo", i));
6294
}
6395
indexRandom(true, indexRequests);
64-
EsqlQueryResponse response = EsqlActionIT.run("from test | stats sum(pause_me)", Settings.EMPTY);
65-
assertThat(response.values(), equalTo(List.of(List.of(5000L))));
6696
}
6797

68-
public static class PausableFieldPlugin extends Plugin implements ScriptPlugin {
98+
public static class TestRuntimeFieldPlugin extends Plugin implements ScriptPlugin {
6999
@Override
70100
public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
71101
return new ScriptEngine() {
@@ -82,22 +112,61 @@ public <FactoryType> FactoryType compile(
82112
ScriptContext<FactoryType> context,
83113
Map<String, String> params
84114
) {
85-
return (FactoryType) new LongFieldScript.Factory() {
86-
@Override
87-
public LongFieldScript.LeafFactory newFactory(
88-
String fieldName,
89-
Map<String, Object> params,
90-
SearchLookup searchLookup,
91-
OnScriptError onScriptError
92-
) {
93-
return ctx -> new LongFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
94-
@Override
95-
public void execute() {
96-
emit(1);
97-
}
98-
};
99-
}
100-
};
115+
if (context == LongFieldScript.CONTEXT) {
116+
return (FactoryType) new LongFieldScript.Factory() {
117+
@Override
118+
public LongFieldScript.LeafFactory newFactory(
119+
String fieldName,
120+
Map<String, Object> params,
121+
SearchLookup searchLookup,
122+
OnScriptError onScriptError
123+
) {
124+
return ctx -> new LongFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
125+
@Override
126+
public void execute() {
127+
emit(1);
128+
}
129+
};
130+
}
131+
};
132+
}
133+
if (context == DoubleFieldScript.CONTEXT) {
134+
return (FactoryType) new DoubleFieldScript.Factory() {
135+
@Override
136+
public DoubleFieldScript.LeafFactory newFactory(
137+
String fieldName,
138+
Map<String, Object> params,
139+
SearchLookup searchLookup,
140+
OnScriptError onScriptError
141+
) {
142+
return ctx -> new DoubleFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
143+
@Override
144+
public void execute() {
145+
emit(1.0);
146+
}
147+
};
148+
}
149+
};
150+
}
151+
if (context == StringFieldScript.CONTEXT) {
152+
return (FactoryType) new StringFieldScript.Factory() {
153+
@Override
154+
public StringFieldScript.LeafFactory newFactory(
155+
String fieldName,
156+
Map<String, Object> params,
157+
SearchLookup searchLookup,
158+
OnScriptError onScriptError
159+
) {
160+
return ctx -> new StringFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
161+
@Override
162+
public void execute() {
163+
emit("const");
164+
}
165+
};
166+
}
167+
};
168+
}
169+
throw new IllegalArgumentException("unsupported context " + context);
101170
}
102171

103172
@Override

0 commit comments

Comments
 (0)