Skip to content

ES|QL: Improve field resolution for FORK #128501

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 28, 2025
Merged

ES|QL: Improve field resolution for FORK #128501

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f8f0478
Improve field resolution for FORK
ioanatia May 27, 2025
1474874
Request all fields when FORK is used
ioanatia May 27, 2025
40b9c7e
Merge branch 'main' into fork_field_resolution
ioanatia May 27, 2025
c2fa62f
Remove unused variable
ioanatia May 28, 2025
2e83b97
Merge branch 'main' into fork_field_resolution
ioanatia May 28, 2025
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
6 changes: 0 additions & 6 deletions muted-tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -453,12 +453,6 @@ tests:
- class: org.elasticsearch.packaging.test.DockerTests
method: test121CanUseStackLoggingConfig
issue: https://github.com/elastic/elasticsearch/issues/128165
- class: org.elasticsearch.xpack.esql.session.IndexResolverFieldNamesTests
method: testForkFieldsWithKeepAfterFork
issue: https://github.com/elastic/elasticsearch/issues/128271
- class: org.elasticsearch.xpack.esql.session.IndexResolverFieldNamesTests
method: testForkWithStatsInAllBranches
issue: https://github.com/elastic/elasticsearch/issues/128272
- class: org.elasticsearch.packaging.test.DockerTests
method: test080ConfigurePasswordThroughEnvironmentVariableFile
issue: https://github.com/elastic/elasticsearch/issues/128075
Expand Down
37 changes: 11 additions & 26 deletions x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -584,38 +584,17 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
return result.withFieldNames(IndexResolver.ALL_FIELDS);
}

Holder<Boolean> projectAll = new Holder<>(false);
parsed.forEachExpressionDown(UnresolvedStar.class, us -> {// explicit "*" fields selection
if (projectAll.get()) {
return;
}
projectAll.set(true);
});

if (projectAll.get()) {
// TODO: Improve field resolution for FORK - right now we request all fields
if (parsed.anyMatch(p -> p instanceof Fork)) {
return result.withFieldNames(IndexResolver.ALL_FIELDS);
}

Holder<Boolean> projectAfterFork = new Holder<>(false);
Holder<Boolean> hasFork = new Holder<>(false);

parsed.forEachDown(plan -> {
Holder<Boolean> projectAll = new Holder<>(false);
parsed.forEachExpressionDown(UnresolvedStar.class, us -> {// explicit "*" fields selection
if (projectAll.get()) {
return;
}

if (hasFork.get() == false && shouldCollectReferencedFields(plan, inlinestatsAggs)) {
projectAfterFork.set(true);
}

if (plan instanceof Fork fork && projectAfterFork.get() == false) {
hasFork.set(true);
fork.children().forEach(child -> {
if (child.anyMatch(p -> shouldCollectReferencedFields(p, inlinestatsAggs)) == false) {
projectAll.set(true);
}
});
}
projectAll.set(true);
});

if (projectAll.get()) {
Expand All @@ -637,6 +616,8 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy

boolean[] canRemoveAliases = new boolean[] { true };

PreAnalysisResult initialResult = result;
projectAll.set(false);
Copy link
Member

@carlosdelest carlosdelest May 28, 2025

Choose a reason for hiding this comment

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

Nit - I think this assignment is unnecessary

parsed.forEachDown(p -> {// go over each plan top-down
if (p instanceof RegexExtract re) { // for Grok and Dissect
// keep the inputs needed by Grok/Dissect
Expand Down Expand Up @@ -711,6 +692,10 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
}
});

if (projectAll.get()) {
Copy link
Member

@carlosdelest carlosdelest May 28, 2025

Choose a reason for hiding this comment

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

This can't happen, right? I don't see projectAll being updated in the previous forEachDown() block 🤔

return result.withFieldNames(IndexResolver.ALL_FIELDS);
}

// Add JOIN ON column references afterward to avoid Alias removal
referencesBuilder.addAll(keepJoinRefsBuilder);
// If any JOIN commands need wildcard field-caps calls, persist the index names
Expand Down
54 changes: 15 additions & 39 deletions ...esql/src/test/java/org/elasticsearch/xpack/esql/session/IndexResolverFieldNamesTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1994,21 +1994,21 @@ public void testForkFieldsWithKeepAfterFork() {
(WHERE d > 1000 AND e == "aaa" | EVAL c = a + 200)
| WHERE x > y
| KEEP a, b, c, d, x
""", Set.of("a", "a.*", "c", "c.*", "d", "d.*", "e", "e.*", "x", "x.*", "y", "y.*"));
""", ALL_FIELDS);
Copy link
Contributor Author

@ioanatia ioanatia May 27, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

technically here what we had before was better
what we do now is that if we detect that a FORK branch requires all fields, we don't look further and we end up requiring all fields from field caps.
I think this is fine for now - we can look into improving this later - but at least the current implementation is less error prone than before.

}

public void testForkFieldsWithKeepBeforeFork() {
assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());

assertFieldNames("""
FROM test
| KEEP a, b, c, d, x
| KEEP a, b, c, d, x, y
| WHERE a > 2000
| EVAL b = a + 100
| FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
(WHERE d > 1000 AND e == "aaa" | EVAL c = a + 200)
| WHERE x > y
""", Set.of("a", "a.*", "b", "b.*", "c", "c.*", "d", "d.*", "e", "e.*", "x", "x.*", "y", "y.*"));
""", ALL_FIELDS);
}

public void testForkFieldsWithNoProjection() {
Expand Down Expand Up @@ -2041,41 +2041,17 @@ public void testForkFieldsWithEnrichAndLookupJoins() {
assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
assumeTrue("LOOKUP JOIN available as snapshot only", EsqlCapabilities.Cap.JOIN_LOOKUP_V12.isEnabled());

assertFieldNames(
"""
FROM test
| KEEP a, b, abc, def, z, xyz
| ENRICH enrich_policy ON abc
| EVAL b = a + 100
| LOOKUP JOIN my_lookup_index ON def
| FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
(STATS x = count(*), y=min(z))
| LOOKUP JOIN my_lookup_index ON xyz
| WHERE x > y OR _fork == "fork1"
""",
Set.of(
"x",
"y",
"_fork",
"a",
"c",
"abc",
"b",
"def",
"z",
"xyz",
"def.*",
"_fork.*",
"y.*",
"x.*",
"xyz.*",
"z.*",
"abc.*",
"a.*",
"c.*",
"b.*"
)
);
assertFieldNames("""
FROM test
| KEEP a, b, abc, def, z, xyz
| ENRICH enrich_policy ON abc
| EVAL b = a + 100
| LOOKUP JOIN my_lookup_index ON def
| FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
(STATS x = count(*), y=min(z))
| LOOKUP JOIN my_lookup_index ON xyz
| WHERE x > y OR _fork == "fork1"
""", ALL_FIELDS);
}

public void testForkWithStatsInAllBranches() {
Expand All @@ -2089,7 +2065,7 @@ public void testForkWithStatsInAllBranches() {
(EVAL z = a * b | STATS m = max(z))
(STATS x = count(*), y=min(z))
| WHERE x > y
""", Set.of("a", "a.*", "b", "b.*", "c", "c.*", "z", "z.*"));
""", ALL_FIELDS);
}

public void testForkWithStatsAndWhere() {
Expand Down