diff --git a/docs/changelog/136072.yaml b/docs/changelog/136072.yaml new file mode 100644 index 0000000000000..c4daef203c112 --- /dev/null +++ b/docs/changelog/136072.yaml @@ -0,0 +1,5 @@ +pr: 136072 +summary: Extend time range bucketing attributes to retrievers +area: Search +type: enhancement +issues: [] diff --git a/server/src/main/java/org/elasticsearch/action/search/SearchRequestAttributesExtractor.java b/server/src/main/java/org/elasticsearch/action/search/SearchRequestAttributesExtractor.java index ef7ef49759790..594369f9e43bf 100644 --- a/server/src/main/java/org/elasticsearch/action/search/SearchRequestAttributesExtractor.java +++ b/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 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"), diff --git a/server/src/main/java/org/elasticsearch/index/query/RankDocsQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/RankDocsQueryBuilder.java index 524310c547597..03f7679584344 100644 --- a/server/src/main/java/org/elasticsearch/index/query/RankDocsQueryBuilder.java +++ b/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; diff --git a/server/src/test/java/org/elasticsearch/action/search/SearchRequestAttributesExtractorTests.java b/server/src/test/java/org/elasticsearch/action/search/SearchRequestAttributesExtractorTests.java index cbb684acb39bf..a5995aa0dac2f 100644 --- a/server/src/test/java/org/elasticsearch/action/search/SearchRequestAttributesExtractorTests.java +++ b/server/src/test/java/org/elasticsearch/action/search/SearchRequestAttributesExtractorTests.java @@ -20,12 +20,18 @@ import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder; import org.elasticsearch.search.builder.PointInTimeBuilder; import org.elasticsearch.search.builder.SearchSourceBuilder; +import org.elasticsearch.search.retriever.KnnRetrieverBuilder; +import org.elasticsearch.search.retriever.StandardRetrieverBuilder; +import org.elasticsearch.search.retriever.TestCompoundRetrieverBuilder; import org.elasticsearch.search.sort.FieldSortBuilder; import org.elasticsearch.search.sort.GeoDistanceSortBuilder; import org.elasticsearch.search.sort.ScoreSortBuilder; import org.elasticsearch.search.sort.ScriptSortBuilder; +import org.elasticsearch.search.vectors.KnnSearchBuilder; +import org.elasticsearch.search.vectors.KnnVectorQueryBuilder; import org.elasticsearch.test.ESTestCase; +import java.util.List; import java.util.Map; public class SearchRequestAttributesExtractorTests extends ESTestCase { @@ -151,53 +157,115 @@ private static void assertAttributes( } } - public void testExtractAttributes() { - { - SearchRequest searchRequest = new SearchRequest(); - Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( - searchRequest, - searchRequest.indices() - ); - assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, null); - } - { - SearchRequest searchRequest = new SearchRequest(); - SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); - searchRequest.source(searchSourceBuilder); - searchSourceBuilder.pointInTimeBuilder(new PointInTimeBuilder(BytesArray.EMPTY)); - Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( - searchRequest, - searchRequest.indices() - ); - assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, "pit"); - } - { - SearchRequest searchRequest = new SearchRequest(); - searchRequest.scroll(new TimeValue(randomIntBetween(1, 10))); - Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( - searchRequest, - searchRequest.indices() - ); - assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, "scroll"); + public void testExtractAttributesTargetOnly() { + SearchRequest searchRequest = new SearchRequest(); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, null); + } + + public void testExtractAttributesTopLevelKnn() { + SearchRequest searchRequest = new SearchRequest(); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.knnSearch(List.of(new KnnSearchBuilder("field", new float[] {}, 2, 5, 10f, null, null))); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", true, false, false, null); + } + + public void testExtractAttributesKnnQuery() { + SearchRequest searchRequest = new SearchRequest(); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.query(new KnnVectorQueryBuilder("field", new float[] {}, 2, 5, 10f, null, null)); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", true, false, false, null); + } + + public void testExtractAttributesKnnRetriever() { + SearchRequest searchRequest = new SearchRequest(); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.retriever(new KnnRetrieverBuilder("field", new float[] {}, null, 2, 5, 10f, null, null)); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", true, false, false, null); + } + + public void testExtractAttributesPit() { + SearchRequest searchRequest = new SearchRequest(); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.pointInTimeBuilder(new PointInTimeBuilder(BytesArray.EMPTY)); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, "pit"); + } + + public void testExtractAttributesScroll() { + SearchRequest searchRequest = new SearchRequest(); + searchRequest.scroll(new TimeValue(randomIntBetween(1, 10))); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, "scroll"); + } + + public void testExtractAttributesTimestampSorted() { + SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.sort("@timestamp"); + searchSourceBuilder.query(new RangeQueryBuilder("@timestamp")); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, false, false, null); + } + + public void testExtractAttributesTimestampSortedTimeRangeFilter() { + SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.sort("@timestamp"); + searchSourceBuilder.query(new RangeQueryBuilder("@timestamp").from("2021-11-11")); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); + } + + public void testExtractAttributesTimestampSortedTimeRangeFilterShouldClauses() { + SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.sort("@timestamp"); + int numBool = randomIntBetween(2, 10); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + for (int i = 0; i < numBool; i++) { + BoolQueryBuilder boolQueryBuilderNew = new BoolQueryBuilder(); + boolQueryBuilder.must(boolQueryBuilderNew); + boolQueryBuilder = boolQueryBuilderNew; } - { - SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); - SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); - searchRequest.source(searchSourceBuilder); - searchSourceBuilder.sort("@timestamp"); - searchSourceBuilder.query(new RangeQueryBuilder("@timestamp")); - Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( - searchRequest, - searchRequest.indices() - ); - assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, false, false, null); + boolQueryBuilder.must(new RangeQueryBuilder("@timestamp").from("2021-11-11")); + searchSourceBuilder.query(boolQueryBuilder); + if (randomBoolean()) { + boolQueryBuilder.should(new RangeQueryBuilder("event.ingested").from("2021-11-11")); } + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); + } + + public void testExtractAttributesTimestampSortedTimeRangeFilters() { { SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchRequest.source(searchSourceBuilder); searchSourceBuilder.sort("@timestamp"); - searchSourceBuilder.query(new RangeQueryBuilder("@timestamp").from("2021-11-11")); + int numBool = randomIntBetween(2, 10); + BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + for (int i = 0; i < numBool; i++) { + BoolQueryBuilder boolQueryBuilderNew = new BoolQueryBuilder(); + boolQueryBuilder.filter(boolQueryBuilderNew); + boolQueryBuilder = boolQueryBuilderNew; + } + if (randomBoolean()) { + boolQueryBuilder.should(new RangeQueryBuilder("event.ingested")); + } + + boolQueryBuilder.filter(new RangeQueryBuilder("@timestamp").from("2021-11-11")); + searchSourceBuilder.query(boolQueryBuilder); Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( searchRequest, searchRequest.indices() @@ -209,24 +277,20 @@ public void testExtractAttributes() { SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchRequest.source(searchSourceBuilder); searchSourceBuilder.sort("@timestamp"); - int numBool = randomIntBetween(2, 10); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); - for (int i = 0; i < numBool; i++) { - BoolQueryBuilder boolQueryBuilderNew = new BoolQueryBuilder(); - boolQueryBuilder.must(boolQueryBuilderNew); - boolQueryBuilder = boolQueryBuilderNew; - } boolQueryBuilder.must(new RangeQueryBuilder("@timestamp").from("2021-11-11")); + boolQueryBuilder.must(new RangeQueryBuilder("event.ingested").from("2021-11-11")); + boolQueryBuilder.must(new RangeQueryBuilder(randomAlphaOfLengthBetween(3, 10))); searchSourceBuilder.query(boolQueryBuilder); - if (randomBoolean()) { - boolQueryBuilder.should(new RangeQueryBuilder("event.ingested").from("2021-11-11")); - } Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( searchRequest, searchRequest.indices() ); - assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); + assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, true, null); } + } + + public void testExtractAttributesTimestampSortedRetriever() { { SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); @@ -244,7 +308,7 @@ public void testExtractAttributes() { } boolQueryBuilder.filter(new RangeQueryBuilder("@timestamp").from("2021-11-11")); - searchSourceBuilder.query(boolQueryBuilder); + searchSourceBuilder.retriever(new StandardRetrieverBuilder(boolQueryBuilder)); Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( searchRequest, searchRequest.indices() @@ -256,17 +320,28 @@ public void testExtractAttributes() { SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); searchRequest.source(searchSourceBuilder); searchSourceBuilder.sort("@timestamp"); + TestCompoundRetrieverBuilder compoundRetrieverBuilder = new TestCompoundRetrieverBuilder(10); + searchSourceBuilder.retriever(compoundRetrieverBuilder); + int numCompound = randomIntBetween(2, 10); + for (int i = 0; i < numCompound; i++) { + TestCompoundRetrieverBuilder innerCompoundRetriever = new TestCompoundRetrieverBuilder(10); + compoundRetrieverBuilder.addChild(innerCompoundRetriever); + compoundRetrieverBuilder = innerCompoundRetriever; + } BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); boolQueryBuilder.must(new RangeQueryBuilder("@timestamp").from("2021-11-11")); boolQueryBuilder.must(new RangeQueryBuilder("event.ingested").from("2021-11-11")); boolQueryBuilder.must(new RangeQueryBuilder(randomAlphaOfLengthBetween(3, 10))); - searchSourceBuilder.query(boolQueryBuilder); + compoundRetrieverBuilder.addChild(new StandardRetrieverBuilder(boolQueryBuilder)); Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( searchRequest, searchRequest.indices() ); assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, true, null); } + } + + public void testExtractAttributesTimestampSortedTimeRangeFilterOneShouldClause() { { SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); @@ -287,6 +362,7 @@ public void testExtractAttributes() { searchRequest.source(searchSourceBuilder); searchSourceBuilder.sort("@timestamp"); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); + // range on some other field boolQueryBuilder.should(new RangeQueryBuilder(randomAlphaOfLengthBetween(3, 10)).from("2021-11-11")); searchSourceBuilder.query(boolQueryBuilder); Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( @@ -295,35 +371,31 @@ public void testExtractAttributes() { ); assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, false, false, null); } - { - SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); - SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); - searchRequest.source(searchSourceBuilder); - searchSourceBuilder.sort("@timestamp"); - searchSourceBuilder.query(new ConstantScoreQueryBuilder(new RangeQueryBuilder("@timestamp").from("2021-11-11"))); - Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( - searchRequest, - searchRequest.indices() - ); - assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); - } - { - SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); - SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); - searchRequest.source(searchSourceBuilder); - searchSourceBuilder.sort("@timestamp"); - searchSourceBuilder.query( - new BoostingQueryBuilder(new RangeQueryBuilder("@timestamp").from("2021-11-11"), new MatchAllQueryBuilder()) - ); - Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( - searchRequest, - searchRequest.indices() - ); - assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); - } } - public void testDepthLimit() { + public void testExtractAttributesTimestampSortedTimeRangeFilterConstantScore() { + SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.sort("@timestamp"); + searchSourceBuilder.query(new ConstantScoreQueryBuilder(new RangeQueryBuilder("@timestamp").from("2021-11-11"))); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); + } + + public void testExtractAttributesTimestampSortedTimeRangeFilterBoostingQuery() { + SearchRequest searchRequest = new SearchRequest(randomAlphaOfLengthBetween(3, 10)); + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchRequest.source(searchSourceBuilder); + searchSourceBuilder.sort("@timestamp"); + searchSourceBuilder.query( + new BoostingQueryBuilder(new RangeQueryBuilder("@timestamp").from("2021-11-11"), new MatchAllQueryBuilder()) + ); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes(searchRequest, searchRequest.indices()); + assertAttributes(stringObjectMap, "user", "@timestamp", "hits_only", false, true, false, null); + } + + public void testIntrospectQueryBuilderDepthLimit() { { SearchRequest searchRequest = new SearchRequest("index"); BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder(); @@ -364,6 +436,43 @@ public void testDepthLimit() { } } + public void testIntrospectCompoundRetrieverBuilderDepthLimit() { + { + SearchRequest searchRequest = new SearchRequest("index"); + TestCompoundRetrieverBuilder compoundRetrieverBuilder = new TestCompoundRetrieverBuilder(10); + searchRequest.source(new SearchSourceBuilder().retriever(compoundRetrieverBuilder)); + int depth = randomIntBetween(5, 18); + for (int i = 0; i < depth; i++) { + TestCompoundRetrieverBuilder innerCmpoundRetrieverBuilder = new TestCompoundRetrieverBuilder(10); + compoundRetrieverBuilder.addChild(innerCmpoundRetrieverBuilder); + compoundRetrieverBuilder = innerCmpoundRetrieverBuilder; + } + compoundRetrieverBuilder.addChild(new StandardRetrieverBuilder(new RangeQueryBuilder("@timestamp").from("2021-11-11"))); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( + searchRequest, + searchRequest.indices() + ); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, true, false, null); + } + { + SearchRequest searchRequest = new SearchRequest("index"); + TestCompoundRetrieverBuilder compoundRetrieverBuilder = new TestCompoundRetrieverBuilder(10); + searchRequest.source(new SearchSourceBuilder().retriever(compoundRetrieverBuilder)); + int depth = randomIntBetween(19, 50); + for (int i = 0; i < depth; i++) { + TestCompoundRetrieverBuilder innerCmpoundRetrieverBuilder = new TestCompoundRetrieverBuilder(10); + compoundRetrieverBuilder.addChild(innerCmpoundRetrieverBuilder); + compoundRetrieverBuilder = innerCmpoundRetrieverBuilder; + } + compoundRetrieverBuilder.addChild(new StandardRetrieverBuilder(new RangeQueryBuilder("@timestamp").from("2021-11-11"))); + Map stringObjectMap = SearchRequestAttributesExtractor.extractAttributes( + searchRequest, + searchRequest.indices() + ); + assertAttributes(stringObjectMap, "user", "_score", "hits_only", false, false, false, null); + } + } + public void testIntrospectTimeRange() { long nowInMillis = System.currentTimeMillis(); assertEquals("15_minutes", SearchRequestAttributesExtractor.introspectTimeRange(nowInMillis, nowInMillis)); diff --git a/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/SearchTookTimeTelemetryTests.java b/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/SearchTookTimeTelemetryTests.java index 4323541a9c8af..2b53b0747a115 100644 --- a/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/SearchTookTimeTelemetryTests.java +++ b/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/SearchTookTimeTelemetryTests.java @@ -745,6 +745,65 @@ public void testTimeRangeFilterAllResultsMixedPrecision() { assertEquals("1_hour", attributes.get("time_range_filter_from")); } + public void testStandardRetrieverWithTimeRangeQuery() { + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchSourceBuilder.retriever(new StandardRetrieverBuilder(new RangeQueryBuilder("event.ingested").from("2024-12-01"))); + SearchResponse searchResponse = client().prepareSearch(indexName).setSource(searchSourceBuilder).get(); + try { + assertNoFailures(searchResponse); + assertSearchHits(searchResponse, "2"); + } finally { + searchResponse.decRef(); + } + + List measurements = getTestTelemetryPlugin().getLongHistogramMeasurement(TOOK_DURATION_TOTAL_HISTOGRAM_NAME); + assertEquals(1, measurements.size()); + assertThat(measurements.getFirst().getLong(), Matchers.lessThanOrEqualTo(searchResponse.getTook().millis())); + assertEquals(searchResponse.getTook().millis(), measurements.getLast().getLong()); + for (Measurement measurement : measurements) { + Map attributes = measurement.attributes(); + assertEquals(5, attributes.size()); + assertEquals("user", attributes.get("target")); + assertEquals("hits_only", attributes.get("query_type")); + assertEquals("_score", attributes.get("sort")); + assertEquals("event.ingested", attributes.get("time_range_filter_field")); + assertEquals("older_than_14_days", attributes.get("time_range_filter_from")); + } + } + + public void testCompoundRetrieverWithTimeRangeQuery() { + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchSourceBuilder.retriever( + new RescorerRetrieverBuilder( + new StandardRetrieverBuilder(new RangeQueryBuilder("@timestamp").from("2024-12-01")), + List.of(new QueryRescorerBuilder(new MatchAllQueryBuilder())) + ) + ); + SearchResponse searchResponse = client().prepareSearch(indexName).setSource(searchSourceBuilder).get(); + try { + assertNoFailures(searchResponse); + assertSearchHits(searchResponse, "2"); + } finally { + searchResponse.decRef(); + } + + List measurements = getTestTelemetryPlugin().getLongHistogramMeasurement(TOOK_DURATION_TOTAL_HISTOGRAM_NAME); + // compound retriever does its own search as an async action, whose took time is recorded separately + assertEquals(2, measurements.size()); + assertThat(measurements.getFirst().getLong(), Matchers.lessThan(searchResponse.getTook().millis())); + assertEquals(searchResponse.getTook().millis(), measurements.getLast().getLong()); + for (Measurement measurement : measurements) { + Map attributes = measurement.attributes(); + assertEquals(6, attributes.size()); + assertEquals("user", attributes.get("target")); + assertEquals("hits_only", attributes.get("query_type")); + assertEquals("_score", attributes.get("sort")); + assertEquals("pit", attributes.get("pit_scroll")); + assertEquals("@timestamp", attributes.get("time_range_filter_field")); + assertEquals("older_than_14_days", attributes.get("time_range_filter_from")); + } + } + private void resetMeter() { getTestTelemetryPlugin().resetMeter(); } diff --git a/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/ShardSearchPhaseAPMMetricsTests.java b/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/ShardSearchPhaseAPMMetricsTests.java index 01c47f9ad2801..95091a0a1e98a 100644 --- a/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/ShardSearchPhaseAPMMetricsTests.java +++ b/server/src/test/java/org/elasticsearch/search/TelemetryMetrics/ShardSearchPhaseAPMMetricsTests.java @@ -15,12 +15,17 @@ import org.elasticsearch.common.settings.Settings; import org.elasticsearch.core.TimeValue; import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.MatchAllQueryBuilder; import org.elasticsearch.index.query.RangeQueryBuilder; import org.elasticsearch.indices.ExecutorNames; import org.elasticsearch.indices.SystemIndexDescriptor; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.PluginsService; import org.elasticsearch.plugins.SystemIndexPlugin; +import org.elasticsearch.search.builder.SearchSourceBuilder; +import org.elasticsearch.search.rescore.QueryRescorerBuilder; +import org.elasticsearch.search.retriever.RescorerRetrieverBuilder; +import org.elasticsearch.search.retriever.StandardRetrieverBuilder; import org.elasticsearch.telemetry.Measurement; import org.elasticsearch.telemetry.TestTelemetryPlugin; import org.elasticsearch.test.ESSingleNodeTestCase; @@ -275,22 +280,59 @@ public void testTimeRangeFilterOneResult() { assertSearchHitsWithoutFailures(client().prepareSearch(TestSystemIndexPlugin.INDEX_NAME).setQuery(rangeQueryBuilder), "2"); final List queryMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(QUERY_SEARCH_PHASE_METRIC); assertEquals(1, queryMeasurements.size()); - assertTimeRangeAttributes(queryMeasurements, ".others", true); + assertTimeRangeAttributes(queryMeasurements, ".others", true, false); final List fetchMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(FETCH_SEARCH_PHASE_METRIC); assertEquals(1, fetchMeasurements.size()); - assertTimeRangeAttributes(fetchMeasurements, ".others", true); + assertTimeRangeAttributes(fetchMeasurements, ".others", true, false); } - private static void assertTimeRangeAttributes(List measurements, String target, boolean isSystem) { + public void testTimeRangeFilterRetrieverOneResult() { + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchSourceBuilder.retriever(new StandardRetrieverBuilder(new RangeQueryBuilder("@timestamp").from("2024-12-01"))); + // target the system index because it has one shard, that simplifies testing. Otherwise, only when the two docs end up indexed + // on the same shard do you get the time range as attribute. + assertSearchHitsWithoutFailures(client().prepareSearch(TestSystemIndexPlugin.INDEX_NAME).setSource(searchSourceBuilder), "2"); + final List queryMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(QUERY_SEARCH_PHASE_METRIC); + assertEquals(1, queryMeasurements.size()); + assertTimeRangeAttributes(queryMeasurements, ".others", true, false); + final List fetchMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(FETCH_SEARCH_PHASE_METRIC); + assertEquals(1, fetchMeasurements.size()); + assertTimeRangeAttributes(fetchMeasurements, ".others", true, false); + } + + public void testTimeRangeFilterCompoundRetrieverOneResult() { + SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder(); + searchSourceBuilder.retriever( + new RescorerRetrieverBuilder( + new StandardRetrieverBuilder(new RangeQueryBuilder("@timestamp").from("2024-12-01")), + List.of(new QueryRescorerBuilder(new MatchAllQueryBuilder())) + ) + ); + // target the system index because it has one shard, that simplifies testing. Otherwise, only when the two docs end up indexed + // on the same shard do you get the time range as attribute. + assertSearchHitsWithoutFailures(client().prepareSearch(TestSystemIndexPlugin.INDEX_NAME).setSource(searchSourceBuilder), "2"); + final List queryMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(QUERY_SEARCH_PHASE_METRIC); + // compound retriever does its own search as an async action, whose metrics are recorded separately + assertEquals(2, queryMeasurements.size()); + assertTimeRangeAttributes(queryMeasurements, ".others", true, true); + final List fetchMeasurements = getTestTelemetryPlugin().getLongHistogramMeasurement(FETCH_SEARCH_PHASE_METRIC); + assertEquals(2, fetchMeasurements.size()); + assertTimeRangeAttributes(fetchMeasurements, ".others", true, true); + } + + private static void assertTimeRangeAttributes(List measurements, String target, boolean isSystem, boolean isPit) { for (Measurement measurement : measurements) { Map attributes = measurement.attributes(); - assertEquals(6, attributes.size()); + assertEquals(isPit ? 7 : 6, attributes.size()); assertEquals(target, attributes.get("target")); assertEquals("hits_only", attributes.get("query_type")); assertEquals("_score", attributes.get("sort")); assertEquals("@timestamp", attributes.get("time_range_filter_field")); assertEquals(isSystem, attributes.get(SearchRequestAttributesExtractor.SYSTEM_THREAD_ATTRIBUTE_NAME)); assertEquals("older_than_14_days", attributes.get("time_range_filter_from")); + if (isPit) { + assertEquals("pit", attributes.get("pit_scroll")); + } } }