[Internal] Query - Fixes AOT violations from Query Pipelines and Query Plan

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Aggregate/Aggregators/SingleGroupAggregator.cs

@@ -230,7 +230,7 @@ public override void AddValues(CosmosElement values)
 
             public override string ToString()
             {
-                return Newtonsoft.Json.JsonConvert.SerializeObject(this.aliasToValue);
+                return System.Text.Json.JsonSerializer.Serialize(this.aliasToValue, CosmosSerializerContext.Default.IReadOnlyDictionaryStringNullableAggregateOperator);
             }
 
             private sealed class OrderedCosmosObject : CosmosObject
@@ -287,7 +287,7 @@ public override void WriteTo(IJsonWriter jsonWriter)
         /// 1) aggregate group by values
         /// 2) scalar group by values.
         /// </summary>
-        private abstract class AggregateValue
+        internal abstract class AggregateValue
         {
             public abstract void AddValue(CosmosElement aggregateValue);
 
@@ -313,7 +313,7 @@ public static TryCatch<AggregateValue> TryCreate(AggregateOperator? aggregateOpe
                 return value;
             }
 
-            private sealed class AggregateAggregateValue : AggregateValue
+            internal sealed class AggregateAggregateValue : AggregateValue
             {
                 private readonly IAggregator aggregator;
 
@@ -374,7 +374,7 @@ public static TryCatch<AggregateValue> TryCreate(
                 }
             }
 
-            private sealed class ScalarAggregateValue : AggregateValue
+            internal sealed class ScalarAggregateValue : AggregateValue
             {
                 private CosmosElement value;
                 private bool initialized;

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CrossPartition/OrderBy/CosmosElementToQueryLiteral.cs

@@ -7,10 +7,10 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy
     using System;
     using System.Collections.Generic;
     using System.Text;
+    using System.Text.Json;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.CosmosElements.Numbers;
     using Microsoft.Azure.Documents.Routing;
-    using Newtonsoft.Json;
 
     internal sealed class CosmosElementToQueryLiteral : ICosmosElementVisitor
     {
@@ -93,7 +93,9 @@ public void Visit(CosmosObject cosmosObject)
 
         public void Visit(CosmosString cosmosString)
         {
-            this.stringBuilder.Append(JsonConvert.SerializeObject(cosmosString.Value.ToString(), DefaultJsonSerializationSettings.Value));
+            this.stringBuilder.Append(JsonSerializer.Serialize(
+                cosmosString.Value.ToString(),
+                CosmosSerializerContext.Default.UtfAnyString));
         }
 
         private sealed class CosmosNumberToQueryLiteral : ICosmosNumberVisitor

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Distinct/DistinctQueryPipelineStage.cs

@@ -6,15 +6,15 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct
 {
     using System;
     using System.Collections.Generic;
+    using System.Text.Json;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.Query.Core.Exceptions;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination;
-    using Microsoft.Azure.Cosmos.Tracing;
-    using Newtonsoft.Json;
+    using Microsoft.Azure.Cosmos.Tracing;
     using UInt128 = Microsoft.Azure.Cosmos.UInt128;
 
     internal class DistinctQueryPipelineStage : QueryPipelineStageBase
@@ -195,7 +195,7 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
         /// <summary>
         /// Continuation token for distinct queries.
         /// </summary>
-        private sealed class DistinctContinuationToken
+        internal sealed class DistinctContinuationToken
         {
             private static class PropertyNames
             {
@@ -255,7 +255,7 @@ public static bool TryParse(
             /// <returns>The serialized form of DistinctContinuationToken</returns>
             public override string ToString()
             {
-                return JsonConvert.SerializeObject(this);
+                return JsonSerializer.Serialize(this, CosmosSerializerContext.Default.DistinctContinuationToken);
             }
         }
     }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Skip/SkipQueryPipelineStage.Client.cs

@@ -8,18 +8,19 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip
     using System.Collections.Generic;
     using System.Diagnostics;
     using System.Linq;
+    using System.Text.Json;
+    using System.Text.Json.Serialization;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.Query.Core.Exceptions;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination;
     using Microsoft.Azure.Cosmos.Tracing;
-    using Newtonsoft.Json;
 
     internal abstract partial class SkipQueryPipelineStage : QueryPipelineStageBase
     {
-        private sealed class ClientSkipQueryPipelineStage : SkipQueryPipelineStage
+        internal sealed class ClientSkipQueryPipelineStage : SkipQueryPipelineStage
         {
             private ClientSkipQueryPipelineStage(
                 IQueryPipelineStage source,
@@ -156,7 +157,7 @@ public override async ValueTask<bool> MoveNextAsync(ITrace trace, CancellationTo
             /// <summary>
             /// A OffsetContinuationToken is a composition of a source continuation token and how many items to skip from that source.
             /// </summary>
-            private readonly struct OffsetContinuationToken
+            internal readonly struct OffsetContinuationToken
             {
                 /// <summary>
                 /// Initializes a new instance of the OffsetContinuationToken struct.
@@ -177,13 +178,13 @@ public OffsetContinuationToken(int offset, string sourceToken)
                 /// <summary>
                 /// The number of items to skip in the query.
                 /// </summary>
-                [JsonProperty("offset")]
+                [JsonPropertyName("offset")]
                 public int Offset { get; }
 
                 /// <summary>
                 /// Gets the continuation token for the source component of the query.
                 /// </summary>
-                [JsonProperty("sourceToken")]
+                [JsonPropertyName("sourceToken")]
                 public string SourceToken { get; }
 
                 /// <summary>
@@ -202,7 +203,9 @@ public static bool TryParse(string value, out OffsetContinuationToken offsetCont
 
                     try
                     {
-                        offsetContinuationToken = JsonConvert.DeserializeObject<OffsetContinuationToken>(value);
+                        offsetContinuationToken = JsonSerializer.Deserialize(
+                            value,
+                            CosmosSerializerContext.Default.OffsetContinuationToken);
                         return true;
                     }
                     catch (JsonException)
@@ -217,7 +220,7 @@ public static bool TryParse(string value, out OffsetContinuationToken offsetCont
                 /// <returns>The string version of the continuation token that can be passed in a response header.</returns>
                 public override string ToString()
                 {
-                    return JsonConvert.SerializeObject(this);
+                    return JsonSerializer.Serialize(this, CosmosSerializerContext.Default.OffsetContinuationToken);
                 }
             }
         }

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/Take/TakeQueryPipelineStage.Client.cs

@@ -7,18 +7,19 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take
     using System;
     using System.Collections.Generic;
     using System.Linq;
+    using System.Text.Json;
+    using System.Text.Json.Serialization;
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.Query.Core.Exceptions;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination;
     using Microsoft.Azure.Cosmos.Tracing;
-    using Newtonsoft.Json;
 
     internal abstract partial class TakeQueryPipelineStage : QueryPipelineStageBase
     {
-        private sealed class ClientTakeQueryPipelineStage : TakeQueryPipelineStage
+        internal sealed class ClientTakeQueryPipelineStage : TakeQueryPipelineStage
         {
             private readonly TakeEnum takeEnum;
 
@@ -235,20 +236,20 @@ public override async ValueTask<bool> MoveNextAsync(ITrace trace, CancellationTo
                 return true;
             }
 
-            private enum TakeEnum
+            internal enum TakeEnum
             {
                 Limit,
                 Top
             }
 
-            private abstract class TakeContinuationToken
+            internal abstract class TakeContinuationToken
             {
             }
 
             /// <summary>
             /// A LimitContinuationToken is a composition of a source continuation token and how many items we have left to drain from that source.
             /// </summary>
-            private sealed class LimitContinuationToken : TakeContinuationToken
+            internal sealed class LimitContinuationToken : TakeContinuationToken
             {
                 /// <summary>
                 /// Initializes a new instance of the LimitContinuationToken struct.
@@ -269,7 +270,7 @@ public LimitContinuationToken(int limit, string sourceToken)
                 /// <summary>
                 /// Gets the limit to the number of document drained for the remainder of the query.
                 /// </summary>
-                [JsonProperty("limit")]
+                [JsonPropertyName("limit")]
                 public int Limit
                 {
                     get;
@@ -278,7 +279,7 @@ public int Limit
                 /// <summary>
                 /// Gets the continuation token for the source component of the query.
                 /// </summary>
-                [JsonProperty("sourceToken")]
+                [JsonPropertyName("sourceToken")]
                 public string SourceToken
                 {
                     get;
@@ -300,7 +301,9 @@ public static bool TryParse(string value, out LimitContinuationToken limitContin
 
                     try
                     {
-                        limitContinuationToken = JsonConvert.DeserializeObject<LimitContinuationToken>(value);
+                        limitContinuationToken = JsonSerializer.Deserialize(
+                            value,
+                            CosmosSerializerContext.Default.LimitContinuationToken);
                         return true;
                     }
                     catch (JsonException)
@@ -315,14 +318,14 @@ public static bool TryParse(string value, out LimitContinuationToken limitContin
                 /// <returns>The string version of the continuation token that can be passed in a response header.</returns>
                 public override string ToString()
                 {
-                    return JsonConvert.SerializeObject(this);
+                    return JsonSerializer.Serialize(this, CosmosSerializerContext.Default.LimitContinuationToken);
                 }
             }
 
             /// <summary>
             /// A TopContinuationToken is a composition of a source continuation token and how many items we have left to drain from that source.
             /// </summary>
-            private sealed class TopContinuationToken : TakeContinuationToken
+            internal sealed class TopContinuationToken : TakeContinuationToken
             {
                 /// <summary>
                 /// Initializes a new instance of the TopContinuationToken struct.
@@ -338,7 +341,7 @@ public TopContinuationToken(int top, string sourceToken)
                 /// <summary>
                 /// Gets the limit to the number of document drained for the remainder of the query.
                 /// </summary>
-                [JsonProperty("top")]
+                [JsonPropertyName("top")]
                 public int Top
                 {
                     get;
@@ -347,7 +350,7 @@ public int Top
                 /// <summary>
                 /// Gets the continuation token for the source component of the query.
                 /// </summary>
-                [JsonProperty("sourceToken")]
+                [JsonPropertyName("sourceToken")]
                 public string SourceToken
                 {
                     get;
@@ -369,7 +372,9 @@ public static bool TryParse(string value, out TopContinuationToken topContinuati
 
                     try
                     {
-                        topContinuationToken = JsonConvert.DeserializeObject<TopContinuationToken>(value);
+                        topContinuationToken = JsonSerializer.Deserialize(
+                            value,
+                            CosmosSerializerContext.Default.TopContinuationToken);
                         return true;
                     }
                     catch (JsonException)
@@ -384,7 +389,7 @@ public static bool TryParse(string value, out TopContinuationToken topContinuati
                 /// <returns>The string version of the continuation token that can be passed in a response header.</returns>
                 public override string ToString()
                 {
-                    return JsonConvert.SerializeObject(this);
+                    return JsonSerializer.Serialize(this, CosmosSerializerContext.Default.TopContinuationToken);
                 }
             }
         }

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/PartitionedQueryExecutionInfo.cs

@@ -61,7 +61,9 @@ public static bool TryParse(string serializedQueryPlan, out PartitionedQueryExec
 
             try
             {
-                partitionedQueryExecutionInfo = JsonSerializer.Deserialize<PartitionedQueryExecutionInfo>(serializedQueryPlan);
+                partitionedQueryExecutionInfo = JsonSerializer.Deserialize(
+                    serializedQueryPlan,
+                    CosmosSerializerContext.Default.PartitionedQueryExecutionInfo);
                 return true;
             }
             catch (JsonException)

Microsoft.Azure.Cosmos/src/Resource/Settings/CosmosSerializerContext.cs

@@ -9,7 +9,8 @@ namespace Microsoft.Azure.Cosmos
     using System.Collections.ObjectModel;
     using System.Text.Json;
     using System.Text.Json.Serialization;
-    using System.Text.Json.Serialization.Metadata;
+    using System.Text.Json.Serialization.Metadata;
+    using Microsoft.Azure.Cosmos.Core.Utf8;
     using Microsoft.Azure.Cosmos.Query.Core;
     using Microsoft.Azure.Cosmos.Query.Core.Metrics;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.OrderBy;
@@ -45,8 +46,13 @@ namespace Microsoft.Azure.Cosmos
     [JsonSerializable(typeof(HybridSearchQueryInfo))]
     [JsonSerializable(typeof(IndexingPolicy))]
     [JsonSerializable(typeof(IndexUtilizationInfo))]
+    [JsonSerializable(typeof(IReadOnlyDictionary<string, Query.Core.Pipeline.Aggregate.Aggregators.SingleGroupAggregator.AggregateValue>))]
     [JsonSerializable(typeof(IReadOnlyList<SortOrder>))]
     [JsonSerializable(typeof(List<CompositeContinuationToken>))]
+    [JsonSerializable(typeof(Microsoft.Azure.Cosmos.Query.Core.Pipeline.Distinct.DistinctQueryPipelineStage.DistinctContinuationToken))]
+    [JsonSerializable(typeof(Microsoft.Azure.Cosmos.Query.Core.Pipeline.Skip.SkipQueryPipelineStage.ClientSkipQueryPipelineStage.OffsetContinuationToken))]
+    [JsonSerializable(typeof(Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take.TakeQueryPipelineStage.ClientTakeQueryPipelineStage.LimitContinuationToken))]
+    [JsonSerializable(typeof(Microsoft.Azure.Cosmos.Query.Core.Pipeline.Take.TakeQueryPipelineStage.ClientTakeQueryPipelineStage.TopContinuationToken))]
     [JsonSerializable(typeof(OtherConnectionConfig))]
     [JsonSerializable(typeof(PartitionKeyDefinition))]
     [JsonSerializable(typeof(PartitionKeyRange))]
@@ -63,6 +69,7 @@ namespace Microsoft.Azure.Cosmos
     [JsonSerializable(typeof(TriggerProperties))]
     [JsonSerializable(typeof(UniqueKeyPolicy))]
     [JsonSerializable(typeof(UserDefinedFunctionProperties))]
+    [JsonSerializable(typeof(UtfAnyString))]
     [JsonSerializable(typeof(VectorEmbeddingPolicy))]
     [JsonSerializable(typeof(HashIndex))]
     [JsonSerializable(typeof(RangeIndex))]