ESQL: Keep DROP attributes when resolving field names (elastic#127009) (elastic#130456)

* ESQL: Keep `DROP` attributes when resolving field names (elastic#127009) (elastic#128147)

(cherry picked from commit 54f2668)

Co-authored-by: kanoshiou <[email protected]>

* [CI] Auto commit changes from spotless

---------

Co-authored-by: kanoshiou <[email protected]>
Co-authored-by: elasticsearchmachine <[email protected]>

Author: 3 people

docs/changelog/127009.yaml

@@ -0,0 +1,6 @@
+pr: 127009
+summary: "ESQL: Keep `DROP` attributes when resolving field names"
+area: ES|QL
+type: bug
+issues:
+  - 126418

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -50,7 +50,6 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
         "token recognition error at: '``", // https://github.com/elastic/elasticsearch/issues/125870
         "Unknown column \\[.*\\]", // https://github.com/elastic/elasticsearch/issues/126026
         "optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/116781
-        "No matches found for pattern", // https://github.com/elastic/elasticsearch/issues/126418
         "The incoming YAML document exceeds the limit:" // still to investigate, but it seems to be specific to the test framework
     );
 

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

@@ -172,3 +172,19 @@ Milky Way
 Milky Way
 Milky Way
 ;
+
+dropAgainWithWildcardAfterEval
+required_capability: drop_again_with_wildcard_after_eval
+from languages
+| eval language_code = 12, x = 13
+| drop language_code
+| drop language*
+| keep x
+| limit 3
+;
+
+x:integer
+13
+13
+13
+;

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

@@ -1567,6 +1567,29 @@ Spanish
 German
 ;
 
+dropAgainWithWildcardAfterEval2
+required_capability: join_lookup_v12
+required_capability: drop_again_with_wildcard_after_eval
+from addresses,cartesian_multipolygons,hosts
+| rename city.name as message 
+| lookup join message_types_lookup on message 
+| eval  card = -6013949614291505456, hOntTwnVC = null, PQAF = null, DXkxCFXyw = null, number = -7336429038807752405 
+| eval  dewAwHC = -1186293612, message = null 
+| sort number ASC, street ASC, ip0 DESC, name ASC NULLS FIRST, host ASC 
+| drop number, host_group, *umber, `city.country.continent.name`, dewAwHC, `zip_code`, `message`, city.country.continent.planet.name, `name`, `ip1`, message, zip_code 
+| drop description, *e, id 
+| keep `hOntTwnVC`, city.country.continent.planet.galaxy, street
+| limit 5
+;
+
+hOntTwnVC:null  | city.country.continent.planet.galaxy:keyword  | street:keyword
+null            | Milky Way                                     | Kearny St
+null            | Milky Way                                     | Keizersgracht
+null            | Milky Way                                     | Marunouchi
+null            | null                                          | null
+null            | null                                          | null
+;
+
 ###############################################
 # LOOKUP JOIN on date_nanos field
 ###############################################

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -734,7 +734,13 @@ public enum Cap {
          * During resolution (pre-analysis) we have to consider that joins or enriches can override EVALuated values
          * https://github.com/elastic/elasticsearch/issues/126419
          */
-        FIX_JOIN_MASKING_EVAL;
+        FIX_JOIN_MASKING_EVAL,
+
+        /**
+         * Support for keeping `DROP` attributes when resolving field names.
+         * see <a href="https://github.com/elastic/elasticsearch/issues/126418"> ES|QL: no matches for pattern #126418 </a>
+         */
+        DROP_AGAIN_WITH_WILDCARD_AFTER_EVAL;
 
         private final boolean enabled;
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -576,10 +576,16 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
         }
 
         AttributeSet references = new AttributeSet();
-        // "keep" attributes are special whenever a wildcard is used in their name
+        // "keep" and "drop" attributes are special whenever a wildcard is used in their name, as the wildcard can shadow some
+        // attributes ("lookup join" generated columns among others) and steps like removal of Aliases should ignore the fields
+        // to remove if their name matches one of these wildcards.
+        //
         // ie "from test | eval lang = languages + 1 | keep *l" should consider both "languages" and "*l" as valid fields to ask for
-        AttributeSet keepCommandReferences = new AttributeSet();
-        AttributeSet keepJoinReferences = new AttributeSet();
+        // "from test | eval first_name = 1 | drop first_name | drop *name should also consider "*name" as valid field to ask for
+        //
+        // NOTE: the grammar allows wildcards to be used in other commands as well, but these are forbidden in the LogicalPlanBuilder
+        var shadowingRefs = new AttributeSet();
+        var keepJoinRefs = new AttributeSet();
         Set<String> wildcardJoinIndices = new java.util.HashSet<>();
 
         boolean[] canRemoveAliases = new boolean[] { true };
@@ -601,14 +607,14 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
                 references.addAll(enrichRefs);
             } else if (p instanceof LookupJoin join) {
                 if (join.config().type() instanceof JoinTypes.UsingJoinType usingJoinType) {
-                    keepJoinReferences.addAll(usingJoinType.columns());
+                    keepJoinRefs.addAll(usingJoinType.columns());
                 }
-                if (keepCommandReferences.isEmpty()) {
+                if (shadowingRefs.isEmpty()) {
                     // No KEEP commands after the JOIN, so we need to mark this index for "*" field resolution
                     wildcardJoinIndices.add(((UnresolvedRelation) join.right()).indexPattern().indexPattern());
                 } else {
                     // Keep commands can reference the join columns with names that shadow aliases, so we block their removal
-                    keepJoinReferences.addAll(keepCommandReferences);
+                    keepJoinRefs.addAll(shadowingRefs);
                 }
             } else {
                 references.addAll(p.references());
@@ -620,12 +626,10 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
                 p.forEachExpression(UnresolvedNamePattern.class, up -> {
                     var ua = new UnresolvedAttribute(up.source(), up.name());
                     references.add(ua);
-                    if (p instanceof Keep) {
-                        keepCommandReferences.add(ua);
-                    }
+                    shadowingRefs.add(ua);
                 });
                 if (p instanceof Keep) {
-                    keepCommandReferences.addAll(p.references());
+                    shadowingRefs.addAll(p.references());
                 }
             }
 
@@ -655,13 +659,13 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
                     if (fieldNames.contains(alias.name())) {
                         return;
                     }
-                    references.removeIf(attr -> matchByName(attr, alias.name(), keepCommandReferences.contains(attr)));
+                    references.removeIf(attr -> matchByName(attr, alias.name(), shadowingRefs.contains(attr)));
                 });
             }
         });
 
         // Add JOIN ON column references afterward to avoid Alias removal
-        references.addAll(keepJoinReferences);
+        references.addAll(keepJoinRefs);
         // If any JOIN commands need wildcard field-caps calls, persist the index names
         if (wildcardJoinIndices.isEmpty() == false) {
             result = result.withWildcardJoinIndices(wildcardJoinIndices);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/session/IndexResolverFieldNamesTests.java

@@ -1632,6 +1632,45 @@ public void testEnrichAndJoinMaskingEvalWh() {
             | keep emp_no, language_name""", Set.of("emp_no", "language_name", "languages", "language_name.*", "languages.*", "emp_no.*"));
     }
 
+    public void testDropAgainWithWildcardAfterEval() {
+        assertFieldNames("""
+            from employees
+            | eval full_name = 12
+            | drop full_name
+            | drop *name
+            | keep emp_no
+            """, Set.of("emp_no", "emp_no.*", "*name", "*name.*"));
+    }
+
+    public void testDropWildcardedFields_AfterRename() {
+        assertFieldNames(
+            """
+                from employees
+                | rename first_name AS first_names, last_name AS last_names
+                | eval first_names = 1
+                | drop first_names
+                | drop *_names
+                | keep gender""",
+            Set.of("first_name", "first_name.*", "last_name", "last_name.*", "*_names", "*_names.*", "gender", "gender.*")
+        );
+    }
+
+    public void testDropWildcardFields_WithLookupJoin() {
+        assumeTrue("LOOKUP JOIN available as snapshot only", EsqlCapabilities.Cap.JOIN_LOOKUP_V12.isEnabled());
+        assertFieldNames(
+            """
+                FROM sample_data
+                | EVAL client_ip = client_ip::keyword
+                | LOOKUP JOIN clientips_lookup ON client_ip
+                | LOOKUP JOIN message_types_lookup ON message
+                | KEEP @timestamp, message, *e*
+                | SORT @timestamp
+                | DROP *e""",
+            Set.of("client_ip", "client_ip.*", "message", "message.*", "@timestamp", "@timestamp.*", "*e*", "*e", "*e.*"),
+            Set.of()
+        );
+    }
+
     private Set<String> fieldNames(String query, Set<String> enrichPolicyMatchFields) {
         var preAnalysisResult = new EsqlSession.PreAnalysisResult(null);
         return EsqlSession.fieldNames(parser.createStatement(query), enrichPolicyMatchFields, preAnalysisResult).fieldNames();