Skip to content

Commit e8a314e

Browse files
dnhatndnhatn
committed
Introduce allow_partial_results setting in ES|QL (elastic#122890)
This change introduces a cluster setting `esql.query.allow_partial_results` that allows enabling or disabling allow_partial_results in ES|QL at the cluster-wide level. Initially, this setting defaults to false, but it will be switched to true soon. The reason for not changing the default in this PR is that it requires adjusting many tests, which would make the PR too large. Instead, we will adjust the tests incrementally and switch the default when the tests are ready. This cluster setting is useful for falling back to the previous behavior (i.e., disabling allow_partial_results) if users upgrade to the new version and haven't updated their queries. Also, the default setting can be overridden on a per-request basis via a URL parameter (allow_partial_results) (changed from request body to URL parameter to conform to the proposal). Relates elastic#122802
1 parent d6512b3 commit e8a314e

File tree

13 files changed

+286
-7
lines changed

13 files changed

+286
-7
lines changed

Diff for: ‎docs/changelog/122890.yaml

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 122890
2+
summary: Introduce `allow_partial_results` setting in ES|QL
3+
area: ES|QL
4+
type: enhancement
5+
issues: []

Diff for: ‎test/external-modules/error-query/build.gradle

+6
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,12 @@
99

1010
import org.elasticsearch.gradle.internal.info.BuildParams
1111
apply plugin: 'elasticsearch.legacy-yaml-rest-test'
12+
apply plugin: 'elasticsearch.internal-java-rest-test'
13+
14+
tasks.named('javaRestTest') {
15+
usesDefaultDistribution()
16+
it.onlyIf("snapshot build") { buildParams.snapshotBuild }
17+
}
1218

1319
tasks.named('yamlRestTest').configure {
1420
it.onlyIf("snapshot build") { buildParams.snapshotBuild }

Diff for: ‎test/external-modules/error-query/src/javaRestTest/java/org/elasticsearch/test/esql/EsqlPartialResultsIT.java

+137
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.test.esql;
11+
12+
import org.apache.http.util.EntityUtils;
13+
import org.elasticsearch.client.Request;
14+
import org.elasticsearch.client.Response;
15+
import org.elasticsearch.client.ResponseException;
16+
import org.elasticsearch.common.settings.Settings;
17+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
18+
import org.elasticsearch.test.cluster.local.distribution.DistributionType;
19+
import org.elasticsearch.test.rest.ESRestTestCase;
20+
import org.junit.ClassRule;
21+
22+
import java.util.HashSet;
23+
import java.util.List;
24+
import java.util.Map;
25+
import java.util.Set;
26+
27+
import static org.hamcrest.Matchers.containsString;
28+
import static org.hamcrest.Matchers.equalTo;
29+
import static org.hamcrest.Matchers.lessThanOrEqualTo;
30+
31+
public class EsqlPartialResultsIT extends ESRestTestCase {
32+
@ClassRule
33+
public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
34+
.distribution(DistributionType.DEFAULT)
35+
.module("test-error-query")
36+
.setting("xpack.security.enabled", "false")
37+
.setting("xpack.license.self_generated.type", "trial")
38+
.setting("esql.query.allow_partial_results", "true")
39+
.build();
40+
41+
@Override
42+
protected String getTestRestCluster() {
43+
return cluster.getHttpAddresses();
44+
}
45+
46+
public Set<String> populateIndices() throws Exception {
47+
int nextId = 0;
48+
{
49+
createIndex("failing-index", Settings.EMPTY, """
50+
{
51+
"runtime": {
52+
"fail_me": {
53+
"type": "long",
54+
"script": {
55+
"source": "",
56+
"lang": "failing_field"
57+
}
58+
}
59+
},
60+
"properties": {
61+
"v": {
62+
"type": "long"
63+
}
64+
}
65+
}
66+
""");
67+
int numDocs = between(1, 50);
68+
for (int i = 0; i < numDocs; i++) {
69+
String id = Integer.toString(nextId++);
70+
Request doc = new Request("PUT", "failing-index/_doc/" + id);
71+
doc.setJsonEntity("{\"v\": " + id + "}");
72+
client().performRequest(doc);
73+
}
74+
75+
}
76+
Set<String> okIds = new HashSet<>();
77+
{
78+
createIndex("ok-index", Settings.EMPTY, """
79+
{
80+
"properties": {
81+
"v": {
82+
"type": "long"
83+
}
84+
}
85+
}
86+
""");
87+
int numDocs = between(1, 50);
88+
for (int i = 0; i < numDocs; i++) {
89+
String id = Integer.toString(nextId++);
90+
okIds.add(id);
91+
Request doc = new Request("PUT", "ok-index/_doc/" + id);
92+
doc.setJsonEntity("{\"v\": " + id + "}");
93+
client().performRequest(doc);
94+
}
95+
}
96+
refresh(client(), "failing-index,ok-index");
97+
return okIds;
98+
}
99+
100+
public void testPartialResult() throws Exception {
101+
Set<String> okIds = populateIndices();
102+
String query = """
103+
{
104+
"query": "FROM ok-index,failing-index | LIMIT 100 | KEEP fail_me,v"
105+
}
106+
""";
107+
// allow_partial_results = true
108+
{
109+
Request request = new Request("POST", "/_query");
110+
request.setJsonEntity(query);
111+
if (randomBoolean()) {
112+
request.addParameter("allow_partial_results", "true");
113+
}
114+
Response resp = client().performRequest(request);
115+
Map<String, Object> results = entityAsMap(resp);
116+
assertThat(results.get("is_partial"), equalTo(true));
117+
List<?> columns = (List<?>) results.get("columns");
118+
assertThat(columns, equalTo(List.of(Map.of("name", "fail_me", "type", "long"), Map.of("name", "v", "type", "long"))));
119+
List<?> values = (List<?>) results.get("values");
120+
assertThat(values.size(), lessThanOrEqualTo(okIds.size()));
121+
}
122+
// allow_partial_results = false
123+
{
124+
Request request = new Request("POST", "/_query");
125+
request.setJsonEntity("""
126+
{
127+
"query": "FROM ok-index,failing-index | LIMIT 100"
128+
}
129+
""");
130+
request.addParameter("allow_partial_results", "false");
131+
var error = expectThrows(ResponseException.class, () -> client().performRequest(request));
132+
Response resp = error.getResponse();
133+
assertThat(resp.getStatusLine().getStatusCode(), equalTo(500));
134+
assertThat(EntityUtils.toString(resp.getEntity()), containsString("Accessing failing field"));
135+
}
136+
}
137+
}

Diff for: ‎test/external-modules/error-query/src/main/java/org/elasticsearch/test/errorquery/ErrorQueryPlugin.java

+54-1
Original file line numberDiff line numberDiff line change
@@ -9,21 +9,74 @@
99

1010
package org.elasticsearch.test.errorquery;
1111

12+
import org.elasticsearch.common.settings.Settings;
13+
import org.elasticsearch.index.mapper.OnScriptError;
1214
import org.elasticsearch.plugins.Plugin;
15+
import org.elasticsearch.plugins.ScriptPlugin;
1316
import org.elasticsearch.plugins.SearchPlugin;
17+
import org.elasticsearch.script.LongFieldScript;
18+
import org.elasticsearch.script.ScriptContext;
19+
import org.elasticsearch.script.ScriptEngine;
20+
import org.elasticsearch.search.lookup.SearchLookup;
1421

22+
import java.util.Collection;
1523
import java.util.List;
24+
import java.util.Map;
25+
import java.util.Set;
1626

1727
import static java.util.Collections.singletonList;
1828

1929
/**
2030
* Test plugin that exposes a way to simulate search shard failures and warnings.
2131
*/
22-
public class ErrorQueryPlugin extends Plugin implements SearchPlugin {
32+
public class ErrorQueryPlugin extends Plugin implements SearchPlugin, ScriptPlugin {
2333
public ErrorQueryPlugin() {}
2434

2535
@Override
2636
public List<QuerySpec<?>> getQueries() {
2737
return singletonList(new QuerySpec<>(ErrorQueryBuilder.NAME, ErrorQueryBuilder::new, p -> ErrorQueryBuilder.PARSER.parse(p, null)));
2838
}
39+
40+
public static final String FAILING_FIELD_LANG = "failing_field";
41+
42+
@Override
43+
public ScriptEngine getScriptEngine(Settings settings, Collection<ScriptContext<?>> contexts) {
44+
return new ScriptEngine() {
45+
@Override
46+
public String getType() {
47+
return FAILING_FIELD_LANG;
48+
}
49+
50+
@Override
51+
@SuppressWarnings("unchecked")
52+
public <FactoryType> FactoryType compile(
53+
String name,
54+
String code,
55+
ScriptContext<FactoryType> context,
56+
Map<String, String> params
57+
) {
58+
return (FactoryType) new LongFieldScript.Factory() {
59+
@Override
60+
public LongFieldScript.LeafFactory newFactory(
61+
String fieldName,
62+
Map<String, Object> params,
63+
SearchLookup searchLookup,
64+
OnScriptError onScriptError
65+
) {
66+
return ctx -> new LongFieldScript(fieldName, params, searchLookup, onScriptError, ctx) {
67+
@Override
68+
public void execute() {
69+
throw new IllegalStateException("Accessing failing field");
70+
}
71+
};
72+
}
73+
};
74+
}
75+
76+
@Override
77+
public Set<ScriptContext<?>> getSupportedContexts() {
78+
return Set.of(LongFieldScript.CONTEXT);
79+
}
80+
};
81+
}
2982
}

Diff for: ‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/esql/action/EsqlQueryRequestBuilder.java

+2
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,6 @@ public final ActionType<Response> action() {
3939

4040
public abstract EsqlQueryRequestBuilder<Request, Response> filter(QueryBuilder filter);
4141

42+
public abstract EsqlQueryRequestBuilder<Request, Response> allowPartialResults(boolean allowPartialResults);
43+
4244
}

Diff for: ‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

+9
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ public static class RequestObjectBuilder {
127127
private Boolean includeCCSMetadata = null;
128128

129129
private CheckedConsumer<XContentBuilder, IOException> filter;
130+
private Boolean allPartialResults = null;
130131

131132
public RequestObjectBuilder() throws IOException {
132133
this(randomFrom(XContentType.values()));
@@ -204,6 +205,11 @@ public RequestObjectBuilder filter(CheckedConsumer<XContentBuilder, IOException>
204205
return this;
205206
}
206207

208+
public RequestObjectBuilder allPartialResults(boolean allPartialResults) {
209+
this.allPartialResults = allPartialResults;
210+
return this;
211+
}
212+
207213
public RequestObjectBuilder build() throws IOException {
208214
if (isBuilt == false) {
209215
if (tables != null) {
@@ -1151,6 +1157,9 @@ static Request prepareRequestWithOptions(RequestObjectBuilder requestObject, Mod
11511157
requestObject.build();
11521158
Request request = prepareRequest(mode);
11531159
String mediaType = attachBody(requestObject, request);
1160+
if (requestObject.allPartialResults != null) {
1161+
request.addParameter("allow_partial_results", String.valueOf(requestObject.allPartialResults));
1162+
}
11541163

11551164
RequestOptions.Builder options = request.getOptions().toBuilder();
11561165
options.setWarningsHandler(WarningsHandler.PERMISSIVE); // We assert the warnings ourselves

Diff for: ‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlNodeFailureIT.java

+44
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@
1717
import org.elasticsearch.xcontent.XContentBuilder;
1818
import org.elasticsearch.xcontent.json.JsonXContent;
1919
import org.elasticsearch.xpack.esql.EsqlTestUtils;
20+
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
2021

2122
import java.util.ArrayList;
2223
import java.util.Collection;
2324
import java.util.HashSet;
2425
import java.util.List;
2526
import java.util.Set;
2627

28+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
2729
import static org.hamcrest.Matchers.equalTo;
2830
import static org.hamcrest.Matchers.in;
2931
import static org.hamcrest.Matchers.lessThanOrEqualTo;
@@ -122,4 +124,46 @@ public void testPartialResults() throws Exception {
122124
}
123125
}
124126
}
127+
128+
public void testDefaultPartialResults() throws Exception {
129+
Set<String> okIds = populateIndices();
130+
assertAcked(
131+
client().admin()
132+
.cluster()
133+
.prepareUpdateSettings(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS)
134+
.setPersistentSettings(Settings.builder().put(EsqlPlugin.QUERY_ALLOW_PARTIAL_RESULTS.getKey(), true))
135+
);
136+
try {
137+
// allow_partial_results = default
138+
{
139+
EsqlQueryRequest request = new EsqlQueryRequest();
140+
request.query("FROM fail,ok | LIMIT 100");
141+
request.pragmas(randomPragmas());
142+
if (randomBoolean()) {
143+
request.allowPartialResults(true);
144+
}
145+
try (EsqlQueryResponse resp = run(request)) {
146+
assertTrue(resp.isPartial());
147+
List<List<Object>> rows = EsqlTestUtils.getValuesList(resp);
148+
assertThat(rows.size(), lessThanOrEqualTo(okIds.size()));
149+
}
150+
}
151+
// allow_partial_results = false
152+
{
153+
EsqlQueryRequest request = new EsqlQueryRequest();
154+
request.query("FROM fail,ok | LIMIT 100");
155+
request.pragmas(randomPragmas());
156+
request.allowPartialResults(false);
157+
IllegalStateException e = expectThrows(IllegalStateException.class, () -> run(request).close());
158+
assertThat(e.getMessage(), equalTo("Accessing failing field"));
159+
}
160+
} finally {
161+
assertAcked(
162+
client().admin()
163+
.cluster()
164+
.prepareUpdateSettings(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS)
165+
.setPersistentSettings(Settings.builder().putNull(EsqlPlugin.QUERY_ALLOW_PARTIAL_RESULTS.getKey()))
166+
);
167+
}
168+
}
125169
}

Diff for: ‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequest.java

+4-3
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public class EsqlQueryRequest extends org.elasticsearch.xpack.core.esql.action.E
5252
private boolean keepOnCompletion;
5353
private boolean onSnapshotBuild = Build.current().isSnapshot();
5454
private boolean acceptedPragmaRisks = false;
55-
private boolean allowPartialResults = false;
55+
private Boolean allowPartialResults = null;
5656

5757
/**
5858
* "Tables" provided in the request for use with things like {@code LOOKUP}.
@@ -232,12 +232,13 @@ public Map<String, Map<String, Column>> tables() {
232232
return tables;
233233
}
234234

235-
public boolean allowPartialResults() {
235+
public Boolean allowPartialResults() {
236236
return allowPartialResults;
237237
}
238238

239-
public void allowPartialResults(boolean allowPartialResults) {
239+
public EsqlQueryRequest allowPartialResults(boolean allowPartialResults) {
240240
this.allowPartialResults = allowPartialResults;
241+
return this;
241242
}
242243

243244
@Override

Diff for: ‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequestBuilder.java

+6
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,12 @@ public EsqlQueryRequestBuilder keepOnCompletion(boolean keepOnCompletion) {
6666
return this;
6767
}
6868

69+
@Override
70+
public EsqlQueryRequestBuilder allowPartialResults(boolean allowPartialResults) {
71+
request.allowPartialResults(allowPartialResults);
72+
return this;
73+
}
74+
6975
static { // plumb access from x-pack core
7076
SharedSecrets.setEsqlQueryRequestBuilderAccess(EsqlQueryRequestBuilder::newSyncEsqlQueryRequestBuilder);
7177
}

Diff for: ‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RequestXContent.java

-2
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,6 @@ String fields() {
8585
static final ParseField WAIT_FOR_COMPLETION_TIMEOUT = new ParseField("wait_for_completion_timeout");
8686
static final ParseField KEEP_ALIVE = new ParseField("keep_alive");
8787
static final ParseField KEEP_ON_COMPLETION = new ParseField("keep_on_completion");
88-
static final ParseField ALLOW_PARTIAL_RESULTS = new ParseField("allow_partial_results");
8988

9089
private static final ObjectParser<EsqlQueryRequest, Void> SYNC_PARSER = objectParserSync(EsqlQueryRequest::syncEsqlQueryRequest);
9190
private static final ObjectParser<EsqlQueryRequest, Void> ASYNC_PARSER = objectParserAsync(EsqlQueryRequest::asyncEsqlQueryRequest);
@@ -115,7 +114,6 @@ private static void objectParserCommon(ObjectParser<EsqlQueryRequest, ?> parser)
115114
parser.declareString((request, localeTag) -> request.locale(Locale.forLanguageTag(localeTag)), LOCALE_FIELD);
116115
parser.declareBoolean(EsqlQueryRequest::profile, PROFILE_FIELD);
117116
parser.declareField((p, r, c) -> new ParseTables(r, p).parseTables(), TABLES_FIELD, ObjectParser.ValueType.OBJECT);
118-
parser.declareBoolean(EsqlQueryRequest::allowPartialResults, ALLOW_PARTIAL_RESULTS);
119117
}
120118

121119
private static ObjectParser<EsqlQueryRequest, Void> objectParserSync(Supplier<EsqlQueryRequest> supplier) {

Diff for: ‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RestEsqlQueryAction.java

+4
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
5151
}
5252

5353
protected static RestChannelConsumer restChannelConsumer(EsqlQueryRequest esqlRequest, RestRequest request, NodeClient client) {
54+
final Boolean partialResults = request.paramAsBoolean("allow_partial_results", null);
55+
if (partialResults != null) {
56+
esqlRequest.allowPartialResults(partialResults);
57+
}
5458
LOGGER.debug("Beginning execution of ESQL query.\nQuery string: [{}]", esqlRequest.query());
5559

5660
return channel -> {

0 commit comments

Comments
 (0)