Query: Adds PopulateIndexMetrics request options (#2612)

This PR adds the option for customer to enable PopulateIndexMetrics in query request options. This allows customer to obtain the index utilization of their query. The information is aggregated and automatically showed up in QueryMetrics once this request option is enabled. Please note that enabling this incurs overhead, so please set it only during debugging. It also adds a field called IndexMetrics to FeedResponse to expose the result to users.

Microsoft.Azure.Cosmos.Encryption/src/Custom/DecryptableFeedResponse.cs

@@ -33,6 +33,9 @@ protected DecryptableFeedResponse(
 
         public override CosmosDiagnostics Diagnostics { get; }
 
+#if SDKPROJECTREF
+        public override string IndexMetrics => null;
+#endif
         public override IEnumerator<T> GetEnumerator()
         {
             return this.Resource.GetEnumerator();

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedEstimatorIterator.cs

@@ -372,6 +372,8 @@ public ChangeFeedEstimatorFeedResponse(
 
             public override CosmosDiagnostics Diagnostics => new CosmosTraceDiagnostics(this.Trace);
 
+            public override string IndexMetrics => null;
+
             public override IEnumerator<ChangeFeedProcessorState> GetEnumerator()
             {
                 return this.remainingLeaseWorks.GetEnumerator();
@@ -403,6 +405,8 @@ public ChangeFeedEstimatorEmptyFeedResponse(ITrace trace)
 
             public override CosmosDiagnostics Diagnostics => new CosmosTraceDiagnostics(this.Trace);
 
+            public override string IndexMetrics => string.Empty;
+
             public override IEnumerator<ChangeFeedProcessorState> GetEnumerator()
             {
                 return ChangeFeedEstimatorEmptyFeedResponse.remainingLeaseWorks.GetEnumerator();

Microsoft.Azure.Cosmos/src/Headers/CosmosMessageHeadersInternal.cs

@@ -87,6 +87,12 @@ public virtual string QueryMetrics
             set => this.SetProperty(HttpConstants.HttpHeaders.QueryMetrics, value);
         }
 
+        public virtual string IndexUtilization
+        {
+            get => this.GetValueOrDefault(HttpConstants.HttpHeaders.IndexUtilization);
+            set => this.SetProperty(HttpConstants.HttpHeaders.IndexUtilization, value);
+        }
+
         public virtual string BackendRequestDurationMilliseconds
         {
             get => this.GetValueOrDefault(HttpConstants.HttpHeaders.BackendRequestDurationMilliseconds);

Microsoft.Azure.Cosmos/src/Headers/CosmosQueryResponseMessageHeaders.cs

@@ -72,7 +72,8 @@ internal CosmosQueryResponseMessageHeaders CloneKnownProperties(
                 RetryAfterLiteral = this.RetryAfterLiteral,
                 SubStatusCodeLiteral = this.SubStatusCodeLiteral,
                 ContentType = this.ContentType,
-                QueryMetricsText = QueryMetricsText
+                QueryMetricsText = QueryMetricsText,
+                IndexUtilizationText = IndexUtilizationText
             };
         }
 
@@ -106,7 +107,8 @@ internal static CosmosQueryResponseMessageHeaders ConvertToQueryHeaders(
                 RetryAfterLiteral = sourceHeaders.RetryAfterLiteral,
                 SubStatusCodeLiteral = sourceHeaders.SubStatusCodeLiteral ?? (substatusCode.HasValue ? substatusCode.Value.ToString() : null),
                 ContentType = sourceHeaders.ContentType,
-                QueryMetricsText = sourceHeaders.QueryMetricsText
+                QueryMetricsText = sourceHeaders.QueryMetricsText,
+                IndexUtilizationText = sourceHeaders.IndexUtilizationText
             };
         }
     }

Microsoft.Azure.Cosmos/src/Headers/Headers.cs

@@ -195,6 +195,12 @@ internal virtual string QueryMetricsText
             set => this.CosmosMessageHeaders.QueryMetrics = value;
         }
 
+        internal virtual string IndexUtilizationText
+        {
+            get => this.CosmosMessageHeaders.IndexUtilization;
+            set => this.CosmosMessageHeaders.IndexUtilization = value;
+        }
+
         internal virtual string BackendRequestDurationMilliseconds
         {
             get => this.CosmosMessageHeaders.BackendRequestDurationMilliseconds;

Microsoft.Azure.Cosmos/src/Query/Core/Metrics/IndexMetricWriter.cs

@@ -0,0 +1,133 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+namespace Microsoft.Azure.Cosmos.Query.Core.Metrics
+{
+    using System;
+    using System.Globalization;
+    using System.Linq;
+    using System.Text;
+
+    /// <summary>
+    /// Base class for visiting and serializing a <see cref="QueryMetrics"/>.
+    /// </summary>
+#if INTERNAL
+#pragma warning disable SA1600
+#pragma warning disable CS1591
+    public
+#else
+    internal
+#endif
+    class IndexMetricWriter
+    {
+        private const string IndexUtilizationInfo = "Index Utilization Information";
+        private const string UtilizedSingleIndexes = "Utilized Single Indexes";
+        private const string PotentialSingleIndexes = "Potential Single Indexes";
+        private const string UtilizedCompositeIndexes = "Utilized Composite Indexes";
+        private const string PotentialCompositeIndexes = "Potential Composite Indexes";
+        private const string IndexExpression = "Index Spec";
+        private const string IndexImpactScore = "Index Impact Score";
+
+        private const string IndexUtilizationSeparator = "---";
+
+        private readonly StringBuilder stringBuilder;
+
+        public IndexMetricWriter(StringBuilder stringBuilder)
+        {
+            this.stringBuilder = stringBuilder ?? throw new ArgumentNullException($"{nameof(stringBuilder)} must not be null.");
+        }
+
+        public void WriteIndexMetrics(IndexUtilizationInfo indexUtilizationInfo)
+        {
+            // IndexUtilizationInfo
+            this.WriteBeforeIndexUtilizationInfo();
+
+            this.WriteIndexUtilizationInfo(indexUtilizationInfo);
+
+            this.WriteAfterIndexUtilizationInfo();
+        }
+
+        #region IndexUtilizationInfo
+        protected void WriteBeforeIndexUtilizationInfo()
+        {
+            IndexMetricWriter.AppendNewlineToStringBuilder(this.stringBuilder);
+            IndexMetricWriter.AppendHeaderToStringBuilder(
+                this.stringBuilder,
+                IndexMetricWriter.IndexUtilizationInfo,
+                indentLevel: 0);
+        }
+
+        protected void WriteIndexUtilizationInfo(IndexUtilizationInfo indexUtilizationInfo)
+        {
+            IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, IndexMetricWriter.UtilizedSingleIndexes, indentLevel: 1);
+
+            foreach (SingleIndexUtilizationEntity indexUtilizationEntity in indexUtilizationInfo.UtilizedSingleIndexes)
+            {
+                WriteSingleIndexUtilizationEntity(indexUtilizationEntity);
+            }
+
+            IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, IndexMetricWriter.PotentialSingleIndexes, indentLevel: 1);
+
+            foreach (SingleIndexUtilizationEntity indexUtilizationEntity in indexUtilizationInfo.PotentialSingleIndexes)
+            {
+                WriteSingleIndexUtilizationEntity(indexUtilizationEntity);
+            }
+
+            IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, IndexMetricWriter.UtilizedCompositeIndexes, indentLevel: 1);
+
+            foreach (CompositeIndexUtilizationEntity indexUtilizationEntity in indexUtilizationInfo.UtilizedCompositeIndexes)
+            {
+                WriteCompositeIndexUtilizationEntity(indexUtilizationEntity);
+            }
+
+            IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, IndexMetricWriter.PotentialCompositeIndexes, indentLevel: 1);
+
+            foreach (CompositeIndexUtilizationEntity indexUtilizationEntity in indexUtilizationInfo.PotentialCompositeIndexes)
+            {
+                WriteCompositeIndexUtilizationEntity(indexUtilizationEntity);
+            }
+
+            void WriteSingleIndexUtilizationEntity(SingleIndexUtilizationEntity indexUtilizationEntity)
+            {
+                IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, $"{IndexMetricWriter.IndexExpression}: {indexUtilizationEntity.IndexDocumentExpression}", indentLevel: 2);
+                IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, $"{IndexMetricWriter.IndexImpactScore}: {indexUtilizationEntity.IndexImpactScore}", indentLevel: 2);
+                IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, IndexMetricWriter.IndexUtilizationSeparator, indentLevel: 2);
+            }
+
+            void WriteCompositeIndexUtilizationEntity(CompositeIndexUtilizationEntity indexUtilizationEntity)
+            {
+                IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, $"{IndexMetricWriter.IndexExpression}: {String.Join(", ", indexUtilizationEntity.IndexDocumentExpressions)}", indentLevel: 2);
+                IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, $"{IndexMetricWriter.IndexImpactScore}: {indexUtilizationEntity.IndexImpactScore}", indentLevel: 2);
+                IndexMetricWriter.AppendHeaderToStringBuilder(this.stringBuilder, IndexMetricWriter.IndexUtilizationSeparator, indentLevel: 2);
+            }
+        }
+
+        protected void WriteAfterIndexUtilizationInfo()
+        {
+            // Do nothing
+        }
+        #endregion
+
+        #region Helpers
+        private static void AppendHeaderToStringBuilder(StringBuilder stringBuilder, string headerTitle, int indentLevel)
+        {
+            const string Indent = "  ";
+            const string FormatString = "{0}{1}";
+
+            stringBuilder.AppendFormat(
+                CultureInfo.InvariantCulture,
+                FormatString,
+                string.Concat(Enumerable.Repeat(Indent, indentLevel)) + headerTitle,
+                Environment.NewLine);
+        }
+
+        private static void AppendNewlineToStringBuilder(StringBuilder stringBuilder)
+        {
+            IndexMetricWriter.AppendHeaderToStringBuilder(
+                stringBuilder,
+                string.Empty,
+                indentLevel: 0);
+        }
+        #endregion
+    }
+}

Microsoft.Azure.Cosmos/src/Query/v3Query/FeedResponse.cs

@@ -43,13 +43,23 @@ protected FeedResponse()
         public abstract int Count { get; }
 
         /// <summary>
-        /// Get an enumerator of the object
+        /// Gets the index utilization metrics to be used for debugging purposes. 
+        /// It's applicable to query response only. Other feed response will return null for this field.
+        /// This result is only available if QueryRequestOptions.PopulateIndexMetrics is set to true.
+        /// </summary>
+        /// <value>
+        /// The index utilization metrics.
+        /// </value>
+        public abstract string IndexMetrics { get; }
+
+        /// <summary>
+        /// Get an enumerator of the object.
         /// </summary>
         /// <returns>An instance of an Enumerator</returns>
         public abstract IEnumerator<T> GetEnumerator();
 
         /// <summary>
-        /// Get an enumerator of the object
+        /// Get an enumerator of the object.
         /// </summary>
         /// <returns>An instance of an Enumerator</returns>
         IEnumerator IEnumerable.GetEnumerator()

Microsoft.Azure.Cosmos/src/Query/v3Query/QueryResponse.cs

@@ -7,7 +7,9 @@ namespace Microsoft.Azure.Cosmos
     using System.Collections.Generic;
     using System.IO;
     using System.Net;
+    using System.Text;
     using Microsoft.Azure.Cosmos.CosmosElements;
+    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
     using Microsoft.Azure.Cosmos.Serializer;
     using Microsoft.Azure.Cosmos.Tracing;
 
@@ -179,6 +181,15 @@ private QueryResponse(
             this.Resource = CosmosElementSerializer.GetResources<T>(
                 cosmosArray: cosmosElements,
                 serializerCore: serializerCore);
+
+            this.IndexUtilizationText = new Lazy<string>(() =>
+            {
+                IndexUtilizationInfo parsedIndexUtilizationInfo = IndexUtilizationInfo.CreateFromString(responseMessageHeaders.IndexUtilizationText);
+                StringBuilder stringBuilder = new StringBuilder();
+                IndexMetricWriter indexMetricWriter = new IndexMetricWriter(stringBuilder);
+                indexMetricWriter.WriteIndexMetrics(parsedIndexUtilizationInfo);
+                return stringBuilder.ToString();
+            });
         }
 
         public override string ContinuationToken => this.Headers.ContinuationToken;
@@ -195,6 +206,10 @@ private QueryResponse(
 
         internal CosmosQueryResponseMessageHeaders QueryHeaders { get; }
 
+        private Lazy<string> IndexUtilizationText { get; }
+
+        public override string IndexMetrics => this.IndexUtilizationText.Value;
+
         public override IEnumerator<T> GetEnumerator()
         {
             return this.Resource.GetEnumerator();

Microsoft.Azure.Cosmos/src/Query/v3Query/ReadFeedResponse.cs

@@ -34,6 +34,8 @@ internal ReadFeedResponse(
 
         public override CosmosDiagnostics Diagnostics { get; }
 
+        public override string IndexMetrics { get; }
+
         public override IEnumerator<T> GetEnumerator()
         {
             return this.Resource.GetEnumerator();

Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs

@@ -92,6 +92,19 @@ public class QueryRequestOptions : RequestOptions
         /// </remarks>
         public PartitionKey? PartitionKey { get; set; }
 
+        /// <summary>
+        /// Gets or sets the <see cref="PopulateIndexMetrics"/> request option for document query requests in the Azure Cosmos DB service.
+        /// </summary>
+        /// <remarks>
+        /// <para> 
+        /// PopulateIndexMetrics is used to obtain the index metrics to understand how the query engine used existing indexes 
+        /// and how it could use potential new indexes.
+        /// The results will be displayed in FeedResponse.IndexMetrics. Please note that this options will incur overhead, so it should be 
+        /// enabled only when debugging slow queries.
+        /// </para>
+        /// </remarks>
+        public bool? PopulateIndexMetrics { get; set; }
+
         /// <summary>
         /// Gets or sets the consistency level required for the request in the Azure Cosmos DB service.
         /// </summary>
@@ -242,6 +255,11 @@ internal override void PopulateRequestOptions(RequestMessage request)
                 request.Headers.Set(HttpConstants.HttpHeaders.EnumerationDirection, this.EnumerationDirection.Value.ToString());
             }
 
+            if (this.PopulateIndexMetrics.HasValue)
+            {
+                request.Headers.Add(HttpConstants.HttpHeaders.PopulateIndexMetrics, this.PopulateIndexMetrics.ToString());
+            }
+
             DedicatedGatewayRequestOptions.PopulateMaxIntegratedCacheStalenessOption(this.DedicatedGatewayRequestOptions, request);
 
             request.Headers.Add(HttpConstants.HttpHeaders.PopulateQueryMetrics, bool.TrueString);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/PopulateIndexMetricsTest.cs

@@ -0,0 +1,79 @@
+namespace Microsoft.Azure.Cosmos.EmulatorTests.Query
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Linq;
+    using System.Text;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.CosmosElements;
+    using Microsoft.Azure.Cosmos.Query.Core.Metrics;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Microsoft.Azure.Documents;
+    using Newtonsoft.Json;
+    using Microsoft.Azure.Cosmos.SDK.EmulatorTests.QueryOracle;
+
+    [TestClass]
+    public sealed class PopulateIndexMetricsTest : QueryTestsBase
+    {
+        [TestMethod]
+        public async Task TestIndexMetricsHeaderExistence()
+        {
+            int seed = (int)(DateTime.UtcNow - new DateTime(1970, 1, 1)).TotalSeconds;
+            uint numberOfDocuments = 1;
+            QueryOracleUtil util = new QueryOracle2(seed);
+            IEnumerable<string> inputDocuments = util.GetDocuments(numberOfDocuments);
+
+            await this.CreateIngestQueryDeleteAsync(
+                ConnectionModes.Direct,
+                CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
+                inputDocuments,
+                ImplementationAsync,
+                "/id");
+
+            static async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject> documents)
+            {
+                string query = string.Format("SELECT * FROM c WHERE c.name = 'ABC' AND c.age > 12");
+
+                // Build the expected string
+                Assert.IsTrue(IndexUtilizationInfo.TryCreateFromDelimitedString("eyJVdGlsaXplZFNpbmdsZUluZGV4ZXMiOlt7IkZpbHRlckV4cHJlc3Npb24iOiIoUk9PVC5uYW1lID0gXCJBQkNcIikiLCJJbmRleFNwZWMiOiJcL25hbWVcLz8iLCJGaWx0ZXJQcmVjaXNlU2V0Ijp0cnVlLCJJbmRleFByZWNpc2VTZXQiOnRydWUsIkluZGV4SW1wYWN0U2NvcmUiOiJIaWdoIn0seyJGaWx0ZXJFeHByZXNzaW9uIjoiKFJPT1QuYWdlID4gMTIpIiwiSW5kZXhTcGVjIjoiXC9hZ2VcLz8iLCJGaWx0ZXJQcmVjaXNlU2V0Ijp0cnVlLCJJbmRleFByZWNpc2VTZXQiOnRydWUsIkluZGV4SW1wYWN0U2NvcmUiOiJIaWdoIn1dLCJQb3RlbnRpYWxTaW5nbGVJbmRleGVzIjpbXSwiVXRpbGl6ZWRDb21wb3NpdGVJbmRleGVzIjpbXSwiUG90ZW50aWFsQ29tcG9zaXRlSW5kZXhlcyI6W3siSW5kZXhTcGVjcyI6WyJcL25hbWUgQVNDIiwiXC9hZ2UgQVNDIl0sIkluZGV4UHJlY2lzZVNldCI6ZmFsc2UsIkluZGV4SW1wYWN0U2NvcmUiOiJIaWdoIn1dfQ==",
+                    out IndexUtilizationInfo parsedInfo));
+                StringBuilder stringBuilder = new StringBuilder();
+                IndexMetricWriter indexMetricWriter = new IndexMetricWriter(stringBuilder);
+                indexMetricWriter.WriteIndexMetrics(parsedInfo);
+                string expectedIndexMetricsString = stringBuilder.ToString();
+
+                // Test using GetItemQueryIterator
+                QueryRequestOptions requestOptions = new QueryRequestOptions() { PopulateIndexMetrics = true };
+                FeedIterator<CosmosElement> itemQuery = container.GetItemQueryIterator<CosmosElement>(
+                    query,
+                    requestOptions: requestOptions);
+
+                while (itemQuery.HasMoreResults)
+                {
+                    FeedResponse<CosmosElement> page = await itemQuery.ReadNextAsync();
+                    Assert.IsTrue(page.Headers.AllKeys().Length > 1);
+                    Assert.IsNotNull(page.Headers.Get(HttpConstants.HttpHeaders.IndexUtilization), "Expected index utilization headers for query");
+                    Assert.IsNotNull(page.IndexMetrics, "Expected index metrics response for query");
+                    Assert.AreEqual(expectedIndexMetricsString, page.IndexMetrics);
+                }
+
+                // Test using Stream API
+                using (FeedIterator feedIterator = container.GetItemQueryStreamIterator(
+                queryText: query,
+                continuationToken: null,
+                requestOptions: new QueryRequestOptions
+                {
+                    PopulateIndexMetrics = true,
+                }))
+                {
+                    using (ResponseMessage response = await feedIterator.ReadNextAsync())
+                    {
+                        Assert.IsNotNull(response.Content);
+                        Assert.IsTrue(response.Headers.AllKeys().Length > 1);
+                        Assert.IsNotNull(response.Headers.Get(HttpConstants.HttpHeaders.IndexUtilization), "Expected index utilization headers for query");
+                    }
+                }
+            }
+        }
+    }
+}