add flag for collecting dimensions

Author: not-napoleon

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

@@ -827,11 +827,17 @@ public record PreAnalysisResult(
         EnrichResolution enrichResolution,
         Set<String> fieldNames,
         Set<String> wildcardJoinIndices,
-        InferenceResolution inferenceResolution
+        InferenceResolution inferenceResolution,
+        boolean collectAllDimensions
     ) {
 
-        public PreAnalysisResult(EnrichResolution enrichResolution, Set<String> fieldNames, Set<String> wildcardJoinIndices) {
-            this(null, new HashMap<>(), enrichResolution, fieldNames, wildcardJoinIndices, InferenceResolution.EMPTY);
+        public PreAnalysisResult(
+            EnrichResolution enrichResolution,
+            Set<String> fieldNames,
+            Set<String> wildcardJoinIndices,
+            boolean collectAllDimensions
+        ) {
+            this(null, new HashMap<>(), enrichResolution, fieldNames, wildcardJoinIndices, InferenceResolution.EMPTY, collectAllDimensions);
         }
 
         PreAnalysisResult withInferenceResolution(InferenceResolution newInferenceResolution) {
@@ -841,7 +847,8 @@ PreAnalysisResult withInferenceResolution(InferenceResolution newInferenceResolu
                 enrichResolution(),
                 fieldNames(),
                 wildcardJoinIndices(),
-                newInferenceResolution
+                newInferenceResolution,
+                collectAllDimensions()
             );
         }
 
@@ -852,7 +859,8 @@ PreAnalysisResult withIndexResolution(IndexResolution newIndexResolution) {
                 enrichResolution(),
                 fieldNames(),
                 wildcardJoinIndices(),
-                inferenceResolution()
+                inferenceResolution(),
+                collectAllDimensions()
             );
         }
 

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

@@ -59,22 +59,22 @@ public class FieldNameUtils {
     public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichResolution enrichResolution) {
 
         // we need the match_fields names from enrich policies and THEN, with an updated list of fields, we call field_caps API
-        var enrichPolicyMatchFields = enrichResolution.resolvedEnrichPolicies()
+        Set<String> enrichPolicyMatchFields = enrichResolution.resolvedEnrichPolicies()
             .stream()
             .map(ResolvedEnrichPolicy::matchField)
             .collect(Collectors.toSet());
 
         // get the field names from the parsed plan combined with the ENRICH match fields from the ENRICH policy
         List<LogicalPlan> inlinestats = parsed.collect(InlineStats.class::isInstance);
         Set<Aggregate> inlinestatsAggs = new HashSet<>();
-        for (var i : inlinestats) {
+        for (LogicalPlan i : inlinestats) {
             inlinestatsAggs.add(((InlineStats) i).aggregate());
         }
 
         if (false == parsed.anyMatch(p -> shouldCollectReferencedFields(p, inlinestatsAggs))) {
             // no explicit columns selection, for example "from employees"
             // also, inlinestats only adds columns to the existent output, its Aggregate shouldn't interfere with potentially using "*"
-            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
+            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of(), false);
         }
 
         Holder<Boolean> projectAll = new Holder<>(false);
@@ -86,7 +86,7 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
         });
 
         if (projectAll.get()) {
-            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
+            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of(), false);
         }
 
         var referencesBuilder = new Holder<>(AttributeSet.builder());
@@ -221,7 +221,7 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
         parsed.forEachDownMayReturnEarly(forEachDownProcessor.get());
 
         if (projectAll.get()) {
-            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
+            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of(), false);
         }
 
         // Add JOIN ON column references afterward to avoid Alias removal
@@ -235,12 +235,12 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
 
         if (fieldNames.isEmpty() && enrichPolicyMatchFields.isEmpty()) {
             // there cannot be an empty list of fields, we'll ask the simplest and lightest one instead: _index
-            return new PreAnalysisResult(enrichResolution, IndexResolver.INDEX_METADATA_FIELD, wildcardJoinIndices);
+            return new PreAnalysisResult(enrichResolution, IndexResolver.INDEX_METADATA_FIELD, wildcardJoinIndices, false);
         } else {
             fieldNames.addAll(subfields(fieldNames));
             fieldNames.addAll(enrichPolicyMatchFields);
             fieldNames.addAll(subfields(enrichPolicyMatchFields));
-            return new PreAnalysisResult(enrichResolution, fieldNames, wildcardJoinIndices);
+            return new PreAnalysisResult(enrichResolution, fieldNames, wildcardJoinIndices, false);
         }
     }