Query: Fixes incorrect FeedResponse.Count when result contains undefined elements

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

@@ -163,6 +163,7 @@ class QueryResponse<T> : FeedResponse<T>
     {
         private readonly CosmosSerializerCore serializerCore;
         private readonly CosmosSerializationFormatOptions serializationOptions;
+        private readonly IReadOnlyList<T> resource;
 
         private QueryResponse(
             HttpStatusCode httpStatusCode,
@@ -178,8 +179,7 @@ private QueryResponse(
             this.serializerCore = serializerCore;
             this.serializationOptions = serializationOptions;
             this.StatusCode = httpStatusCode;
-            this.Count = cosmosElements.Count;
-            this.Resource = CosmosElementSerializer.GetResources<T>(
+            this.resource = CosmosElementSerializer.GetResources<T>(
                 cosmosArray: cosmosElements,
                 serializerCore: serializerCore);
 
@@ -197,7 +197,7 @@ private QueryResponse(
 
         public override CosmosDiagnostics Diagnostics { get; }
 
-        public override int Count { get; }
+        public override int Count => this.resource.Count;
 
         internal CosmosQueryResponseMessageHeaders QueryHeaders { get; }
 
@@ -210,7 +210,7 @@ public override IEnumerator<T> GetEnumerator()
             return this.Resource.GetEnumerator();
         }
 
-        public override IEnumerable<T> Resource { get; }
+        public override IEnumerable<T> Resource => this.resource;
 
         internal override RequestMessage RequestMessage { get; }
 

Microsoft.Azure.Cosmos/src/Serializer/CosmosElementSerializer.cs

@@ -74,8 +74,11 @@ internal static MemoryStream ToStream(
             jsonWriter.WriteArrayStart();
             foreach (CosmosElement element in cosmosElements)
             {
-                count++;
-                element.WriteTo(jsonWriter);
+                if (element is not CosmosUndefined)
+                {
+                    count++;
+                    element.WriteTo(jsonWriter);
+                }
             }
 
             jsonWriter.WriteArrayEnd();

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

@@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query
     using System.Collections.ObjectModel;
     using System.Linq;
     using System.Threading.Tasks;
+    using Azure;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.CosmosElements.Numbers;
     using Microsoft.Azure.Cosmos.Json;
@@ -50,6 +51,70 @@ private static async Task RunTests(Container container, IReadOnlyList<CosmosObje
         {
             await OrderByTests(container);
             await GroupByTests(container);
+            await UntypedTests(container);
+        }
+
+        private static async Task UntypedTests(Container container)
+        {
+            UndefinedProjectionTestCase[] undefinedProjectionTestCases = new[]
+            {
+                MakeUndefinedProjectionTest(
+                    query: "SELECT VALUE c.AlwaysUndefinedField FROM c",
+                    expectedCount: 0),
+                MakeUndefinedProjectionTest(
+                    query: "SELECT VALUE c.AlwaysUndefinedField FROM c ORDER BY c.AlwaysUndefinedField",
+                    expectedCount: 0),
+                MakeUndefinedProjectionTest(
+                    query: "SELECT c.AlwaysUndefinedField FROM c ORDER BY c.AlwaysUndefinedField",
+                    expectedCount: DocumentCount),
+                MakeUndefinedProjectionTest(
+                    query: "SELECT VALUE c.AlwaysUndefinedField FROM c GROUP BY c.AlwaysUndefinedField",
+                    expectedCount: 0),
+                MakeUndefinedProjectionTest(
+                    query: $"SELECT VALUE SUM(c.{nameof(MixedTypeDocument.MixedTypeField)}) FROM c",
+                    expectedCount: 0),
+                MakeUndefinedProjectionTest(
+                    query: $"SELECT VALUE AVG(c.{nameof(MixedTypeDocument.MixedTypeField)}) FROM c",
+                    expectedCount: 0),
+            };
+
+            foreach (UndefinedProjectionTestCase testCase in undefinedProjectionTestCases)
+            {
+                foreach (int pageSize in PageSizes)
+                {
+                    IAsyncEnumerable<ResponseMessage> results = RunSimpleQueryAsync(
+                        container,
+                        testCase.Query,
+                        new QueryRequestOptions { MaxItemCount = pageSize });
+
+                    long actualCount = 0;
+                    await foreach (ResponseMessage responseMessage in results)
+                    {
+                        Assert.IsTrue(responseMessage.IsSuccessStatusCode);
+
+                        string content = responseMessage.Content.ReadAsString();
+                        IJsonNavigator navigator = JsonNavigator.Create(System.Text.Encoding.UTF8.GetBytes(content));
+                        IJsonNavigatorNode rootNode = navigator.GetRootNode();
+                        Assert.IsTrue(navigator.TryGetObjectProperty(rootNode, "_count", out ObjectProperty countProperty));
+
+                        long count = Number64.ToLong(navigator.GetNumber64Value(countProperty.ValueNode));
+                        actualCount += count;
+
+                        Assert.IsTrue(navigator.TryGetObjectProperty(rootNode, "Documents", out ObjectProperty documentsProperty));
+                        int documentCount = navigator.GetArrayItemCount(documentsProperty.ValueNode);
+                        Assert.AreEqual(count, documentCount);
+
+                        for (int index= 0; index < documentCount; ++index)
+                        {
+                            IJsonNavigatorNode documentNode = navigator.GetArrayItemAt(documentsProperty.ValueNode, index);
+                            int propertyCount = navigator.GetObjectPropertyCount(documentNode);
+                            Assert.AreEqual(0, propertyCount);
+                        }
+                    }
+
+                    Assert.AreEqual(testCase.ExpectedResultCount, actualCount);
+                }
+            }
         }
 
         private static async Task OrderByTests(Container container)

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

@@ -645,6 +645,7 @@ internal static async Task<List<T>> QueryWithContinuationTokensAsync<T>(
             }
 
             List<T> resultsFromContinuationToken = new List<T>();
+            int resultCount = 0;
             string continuationToken = null;
             do
             {
@@ -668,6 +669,7 @@ internal static async Task<List<T>> QueryWithContinuationTokensAsync<T>(
                             }
 
                             resultsFromContinuationToken.AddRange(cosmosQueryResponse);
+                            resultCount += cosmosQueryResponse.Count;
                             continuationToken = cosmosQueryResponse.ContinuationToken;
                             break;
                         }
@@ -687,6 +689,7 @@ internal static async Task<List<T>> QueryWithContinuationTokensAsync<T>(
                 }
             } while (continuationToken != null);
 
+            Assert.AreEqual(resultsFromContinuationToken.Count, resultCount);
             return resultsFromContinuationToken;
         }
 
@@ -700,6 +703,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
                 queryRequestOptions = new QueryRequestOptions();
             }
 
+            int resultCount = 0;
             List<T> results = new List<T>();
             FeedIterator<T> itemQuery = container.GetItemQueryIterator<T>(
                 queryText: query,
@@ -713,6 +717,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
                     {
                         FeedResponse<T> page = await itemQuery.ReadNextAsync();
                         results.AddRange(page);
+                        resultCount += page.Count;
 
                         if (queryRequestOptions.MaxItemCount.HasValue)
                         {
@@ -746,6 +751,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
                         {
                             // The query failed and we don't have a save point, so just restart the whole thing.
                             results = new List<T>();
+                            resultCount = 0;
                         }
                     }
                 }
@@ -755,6 +761,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
                 itemQuery.Dispose();
             }
 
+            Assert.AreEqual(results.Count, resultCount);
             return results;
         }
 
@@ -907,6 +914,23 @@ internal static async IAsyncEnumerable<FeedResponse<T>> RunSimpleQueryWithNewIte
             }
         }
 
+        internal static async IAsyncEnumerable<ResponseMessage> RunSimpleQueryAsync(
+           Container container,
+           string query,
+           QueryRequestOptions requestOptions = null)
+        {
+            using FeedIterator resultSetIterator = container.GetItemQueryStreamIterator(
+                query,
+                null,
+                requestOptions: requestOptions);
+
+            while (resultSetIterator.HasMoreResults)
+            {
+                ResponseMessage response = await resultSetIterator.ReadNextAsync();
+                yield return response;
+            }
+        }
+
         internal async Task<List<T>> RunSinglePartitionQuery<T>(
             Container container,
             string query,