Extend time range bucketing attributes to retrievers (#136072)

There's situations where the time range filter is provided as part of the retriever tree.
In that case, we capture the time range filter from while parsing it, but we don't do
the corresponding introspection of the retriever tree to extract which fields were
the time range filters made against.

This commit introduces a very basic introspection of retrievers and tests around it, expanding on #135549.

Author: javanna

docs/changelog/136072.yaml

@@ -0,0 +1,5 @@
+pr: 136072
+summary: Extend time range bucketing attributes to retrievers
+area: Search
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/action/search/SearchRequestAttributesExtractor.java

@@ -22,9 +22,14 @@
 import org.elasticsearch.index.query.NestedQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.RangeQueryBuilder;
+import org.elasticsearch.index.query.RankDocsQueryBuilder;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.ShardSearchRequest;
+import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
+import org.elasticsearch.search.retriever.KnnRetrieverBuilder;
+import org.elasticsearch.search.retriever.RetrieverBuilder;
+import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
 import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.search.sort.ScoreSortBuilder;
 import org.elasticsearch.search.sort.SortBuilder;
@@ -108,7 +113,15 @@ private static Map<String, Object> extractAttributes(
             try {
                 introspectQueryBuilder(searchSourceBuilder.query(), queryMetadataBuilder, 0);
             } catch (Exception e) {
-                logger.error("Failed to extract query attribute", e);
+                logger.error("Failed to extract query attributes", e);
+            }
+        }
+
+        if (searchSourceBuilder.retriever() != null) {
+            try {
+                introspectRetriever(searchSourceBuilder.retriever(), queryMetadataBuilder, 0);
+            } catch (Exception e) {
+                logger.error("Failed to extract retriever attributes", e);
             }
         }
 
@@ -311,6 +324,14 @@ private static void introspectQueryBuilder(QueryBuilder queryBuilder, QueryMetad
             case NestedQueryBuilder nested:
                 introspectQueryBuilder(nested.query(), queryMetadataBuilder, ++level);
                 break;
+            case RankDocsQueryBuilder rankDocs:
+                QueryBuilder[] queryBuilders = rankDocs.getQueryBuilders();
+                if (queryBuilders != null) {
+                    for (QueryBuilder builder : queryBuilders) {
+                        introspectQueryBuilder(builder, queryMetadataBuilder, level + 1);
+                    }
+                }
+                break;
             case RangeQueryBuilder range:
                 // Note that the outcome of this switch differs depending on whether it is executed on the coord node, or data node.
                 // Data nodes perform query rewrite on each shard. That means that a query that reports a certain time range filter at the
@@ -338,6 +359,26 @@ private static void introspectQueryBuilder(QueryBuilder queryBuilder, QueryMetad
         }
     }
 
+    private static void introspectRetriever(RetrieverBuilder retrieverBuilder, QueryMetadataBuilder queryMetadataBuilder, int level) {
+        if (level > 20) {
+            return;
+        }
+        switch (retrieverBuilder) {
+            case KnnRetrieverBuilder knn:
+                queryMetadataBuilder.knnQuery = true;
+                break;
+            case StandardRetrieverBuilder standard:
+                introspectQueryBuilder(standard.topDocsQuery(), queryMetadataBuilder, level + 1);
+                break;
+            case CompoundRetrieverBuilder<?> compound:
+                for (CompoundRetrieverBuilder.RetrieverSource retrieverSource : compound.innerRetrievers()) {
+                    introspectRetriever(retrieverSource.retriever(), queryMetadataBuilder, level + 1);
+                }
+                break;
+            default:
+        }
+    }
+
     private enum TimeRangeBucket {
         FifteenMinutes(TimeValue.timeValueMinutes(15).getMillis(), "15_minutes"),
         OneHour(TimeValue.timeValueHours(1).getMillis(), "1_hour"),

server/src/main/java/org/elasticsearch/index/query/RankDocsQueryBuilder.java

@@ -91,6 +91,10 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         }
     }
 
+    public QueryBuilder[] getQueryBuilders() {
+        return queryBuilders;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;