ESQL load all dimensions when doing a Time Series query (elastic#132687)

Fetch all dimensions when running a time series query

Note on CCS - this is adding a filter to field caps, which will not be available on older versions. This will cause the CCS field caps request to fail for those clusters. However, since the TS command this supports will not be released until after this is merged, that will only happen on clusters that weren't going to be able to run the query anyway. This seems acceptable to me, but worth mentioning.

Author: not-napoleon

server/src/main/java/org/elasticsearch/action/fieldcaps/FieldCapabilitiesFetcher.java

@@ -161,6 +161,7 @@ static Map<String, IndexFieldCapabilities> retrieveFieldCaps(
         boolean includeEmptyFields
     ) {
         boolean includeParentObjects = checkIncludeParents(filters);
+        boolean includeDimensions = checkIncludeDimensions(filters);
 
         Predicate<MappedFieldType> filter = buildFilter(filters, types, context);
         boolean isTimeSeriesIndex = context.getIndexSettings().getTimestampBounds() != null;
@@ -169,10 +170,10 @@ static Map<String, IndexFieldCapabilities> retrieveFieldCaps(
         Map<String, IndexFieldCapabilities> responseMap = new HashMap<>();
         for (Map.Entry<String, MappedFieldType> entry : context.getAllFields()) {
             final String field = entry.getKey();
-            if (fieldNameFilter.test(field) == false) {
+            MappedFieldType ft = entry.getValue();
+            if (fieldNameFilter.test(field) == false && ((includeDimensions && ft.isDimension()) == false)) {
                 continue;
             }
-            MappedFieldType ft = entry.getValue();
             if ((includeEmptyFields || ft.fieldHasValue(fieldInfos))
                 && (fieldPredicate.test(ft.name()) || context.isMetadataField(ft.name()))
                 && (filter == null || filter.test(ft))) {
@@ -234,6 +235,15 @@ private static boolean checkIncludeParents(String[] filters) {
         return true;
     }
 
+    private static boolean checkIncludeDimensions(String[] filters) {
+        for (String filter : filters) {
+            if ("+dimension".equals(filter)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private static boolean canMatchShard(
         ShardId shardId,
         QueryBuilder indexFilter,
@@ -267,7 +277,8 @@ private static Predicate<MappedFieldType> buildFilter(String[] filters, String[]
         }
 
         for (String filter : filters) {
-            if ("parent".equals(filter) || "-parent".equals(filter)) {
+            // These "filters" are handled differently, in that they are not ANDed with the field name pattern
+            if ("parent".equals(filter) || "-parent".equals(filter) || "+dimension".equals(filter)) {
                 continue;
             }
             Predicate<MappedFieldType> next = switch (filter) {

server/src/test/java/org/elasticsearch/action/fieldcaps/FieldCapabilitiesFilterTests.java

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.index.FieldInfos;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.MapperServiceTestCase;
 import org.elasticsearch.index.query.SearchExecutionContext;
@@ -97,6 +98,54 @@ public void testMetadataFilters() throws IOException {
         }
     }
 
+    public void testDimensionFilters() throws IOException {
+        MapperService mapperService = createMapperService(
+            Settings.builder().put("index.mode", "time_series").put("index.routing_path", "dim.*").build(),
+            """
+                { "_doc" : {
+                  "properties" : {
+                    "metric" : { "type" : "long" },
+                    "dimension_1" : { "type" : "keyword", "time_series_dimension" : "true" },
+                    "dimension_2" : { "type" : "long", "time_series_dimension" : "true" }
+                  }
+                } }
+                """
+        );
+        SearchExecutionContext sec = createSearchExecutionContext(mapperService);
+
+        {
+            // First, test without the filter
+            Map<String, IndexFieldCapabilities> response = FieldCapabilitiesFetcher.retrieveFieldCaps(
+                sec,
+                s -> s.equals("metric"),
+                Strings.EMPTY_ARRAY,
+                Strings.EMPTY_ARRAY,
+                FieldPredicate.ACCEPT_ALL,
+                getMockIndexShard(),
+                true
+            );
+            assertNotNull(response.get("metric"));
+            assertNull(response.get("dimension_1"));
+            assertNull(response.get("dimension_2"));
+        }
+
+        {
+            // then, test with the filter
+            Map<String, IndexFieldCapabilities> response = FieldCapabilitiesFetcher.retrieveFieldCaps(
+                sec,
+                s -> s.equals("metric"),
+                new String[] { "+dimension" },
+                Strings.EMPTY_ARRAY,
+                FieldPredicate.ACCEPT_ALL,
+                getMockIndexShard(),
+                true
+            );
+            assertNotNull(response.get("dimension_1"));
+            assertNotNull(response.get("dimension_2"));
+            assertNotNull(response.get("metric"));
+        }
+    }
+
     public void testExcludeMultifields() throws IOException {
         MapperService mapperService = createMapperService("""
             { "_doc" : {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/enrich/EnrichPolicyResolver.java

@@ -449,7 +449,7 @@ public void messageReceived(LookupRequest request, TransportChannel channel, Tas
                             } else {
                                 failures.put(policyName, indexResult.toString());
                             }
-                        }));
+                        }), false);
                     }
                 }
             }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/UnresolvedRelation.java

@@ -166,4 +166,12 @@ public List<Object> nodeProperties() {
     public String toString() {
         return UNRESOLVED_PREFIX + indexPattern.indexPattern();
     }
+
+    /**
+     * @return true if and only if this relation is being loaded in "time series mode",
+     *         which changes a number of behaviors in the planner.
+     */
+    public boolean isTimeSeriesMode() {
+        return indexMode == IndexMode.TIME_SERIES;
+    }
 }

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

@@ -424,7 +424,8 @@ private void preAnalyzeLookupIndex(
             patternWithRemotes,
             fieldNames,
             null,
-            listener.map(indexResolution -> receiveLookupIndexResolution(result, localPattern, executionInfo, indexResolution))
+            listener.map(indexResolution -> receiveLookupIndexResolution(result, localPattern, executionInfo, indexResolution)),
+            false
         );
     }
 
@@ -638,8 +639,10 @@ private void preAnalyzeMainIndices(
                     result.withIndexResolution(IndexResolution.valid(new EsIndex(table.indexPattern(), Map.of(), Map.of())))
                 );
             } else {
+                boolean includeAllDimensions = false;
                 // call the EsqlResolveFieldsAction (field-caps) to resolve indices and get field types
                 if (preAnalysis.indexMode == IndexMode.TIME_SERIES) {
+                    includeAllDimensions = true;
                     // TODO: Maybe if no indices are returned, retry without index mode and provide a clearer error message.
                     var indexModeFilter = new TermQueryBuilder(IndexModeFieldMapper.NAME, IndexMode.TIME_SERIES.getName());
                     if (requestFilter != null) {
@@ -654,7 +657,8 @@ private void preAnalyzeMainIndices(
                     requestFilter,
                     listener.delegateFailure((l, indexResolution) -> {
                         l.onResponse(result.withIndexResolution(indexResolution));
-                    })
+                    }),
+                    includeAllDimensions
                 );
             }
         } else {

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

@@ -8,7 +8,6 @@
 package org.elasticsearch.xpack.esql.session;
 
 import org.elasticsearch.common.regex.Regex;
-import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.xpack.esql.analysis.EnrichResolution;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
@@ -64,15 +63,15 @@ 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());
         }
 
@@ -167,7 +166,7 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
                 }
             } else {
                 referencesBuilder.get().addAll(p.references());
-                if (p instanceof UnresolvedRelation ur && ur.indexMode() == IndexMode.TIME_SERIES) {
+                if (p instanceof UnresolvedRelation ur && ur.isTimeSeriesMode()) {
                     // METRICS aggs generally rely on @timestamp without the user having to mention it.
                     referencesBuilder.get().add(new UnresolvedAttribute(ur.source(), MetadataAttribute.TIMESTAMP_FIELD));
                 }

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

@@ -81,11 +81,12 @@ public void resolveAsMergedMapping(
         String indexWildcard,
         Set<String> fieldNames,
         QueryBuilder requestFilter,
-        ActionListener<IndexResolution> listener
+        ActionListener<IndexResolution> listener,
+        boolean includeAllDimensions
     ) {
         client.execute(
             EsqlResolveFieldsAction.TYPE,
-            createFieldCapsRequest(indexWildcard, fieldNames, requestFilter),
+            createFieldCapsRequest(indexWildcard, fieldNames, requestFilter, includeAllDimensions),
             listener.delegateFailureAndWrap((l, response) -> l.onResponse(mergedMappings(indexWildcard, response)))
         );
     }
@@ -273,7 +274,12 @@ private static EsField conflictingMetricTypes(String name, String fullName, Fiel
         return new InvalidMappedField(name, "mapped as different metric types in indices: " + indices);
     }
 
-    private static FieldCapabilitiesRequest createFieldCapsRequest(String index, Set<String> fieldNames, QueryBuilder requestFilter) {
+    private static FieldCapabilitiesRequest createFieldCapsRequest(
+        String index,
+        Set<String> fieldNames,
+        QueryBuilder requestFilter,
+        boolean includeAllDimensions
+    ) {
         FieldCapabilitiesRequest req = new FieldCapabilitiesRequest().indices(Strings.commaDelimitedListToStringArray(index));
         req.fields(fieldNames.toArray(String[]::new));
         req.includeUnmapped(true);
@@ -282,7 +288,11 @@ private static FieldCapabilitiesRequest createFieldCapsRequest(String index, Set
         // also because this way security doesn't throw authorization exceptions but rather honors ignore_unavailable
         req.indicesOptions(FIELD_CAPS_INDICES_OPTIONS);
         // we ignore the nested data type fields starting with https://github.com/elastic/elasticsearch/pull/111495
-        req.filters("-nested");
+        if (includeAllDimensions) {
+            req.filters("-nested", "+dimension");
+        } else {
+            req.filters("-nested");
+        }
         req.setMergeResults(false);
         return req;
     }