Query: Adds LINQ Support for FirstOrDefault

Microsoft.Azure.Cosmos/src/Linq/CosmosLinqExtensions.cs

@@ -10,7 +10,6 @@ namespace Microsoft.Azure.Cosmos.Linq
  using System.Linq;
  using System.Linq.Expressions;
  using System.Reflection;
- using System.Runtime.InteropServices;
  using System.Threading;
  using System.Threading.Tasks;
  using Microsoft.Azure.Cosmos.Diagnostics;
@@ -774,7 +773,7 @@ public static Task<Response<int>> SumAsync(
  return ResponseHelperAsync(source.Sum());
  }
 
- return ((CosmosLinqQueryProvider)source.Provider).ExecuteAggregateAsync<int?>(
+ return cosmosLinqQueryProvider.ExecuteAggregateAsync<int?>(
  Expression.Call(
  GetMethodInfoOf<IQueryable<int?>, int?>(Queryable.Sum),
  source.Expression),

Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQuery.cs

@@ -16,6 +16,7 @@ namespace Microsoft.Azure.Cosmos.Linq
  using Microsoft.Azure.Cosmos.Serializer;
  using Microsoft.Azure.Cosmos.Tracing;
  using Newtonsoft.Json;
+ using Debug = System.Diagnostics.Debug;
 
  /// <summary>
  /// This is the entry point for LINQ query creation/execution, it generate query provider, implements IOrderedQueryable.
@@ -108,7 +109,12 @@ public IEnumerator<T> GetEnumerator()
  " use GetItemQueryIterator to execute asynchronously");
  }
 
- FeedIterator<T> localFeedIterator = this.CreateFeedIterator(false);
+ FeedIterator<T> localFeedIterator = this.CreateFeedIterator(false, out ScalarOperationKind scalarOperationKind);
+ Debug.Assert(
+ scalarOperationKind == ScalarOperationKind.None,
+ "CosmosLinqQuery Assert!",
+ $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'");
+
  while (localFeedIterator.HasMoreResults)
  {
 #pragma warning disable VSTHRD002 // Avoid problematic synchronous waits
@@ -133,7 +139,7 @@ IEnumerator IEnumerable.GetEnumerator()
 
  public override string ToString()
  {
- SqlQuerySpec querySpec = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions);
+ SqlQuerySpec querySpec = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions).SqlQuerySpec;
  if (querySpec != null)
  {
  return JsonConvert.SerializeObject(querySpec);
@@ -144,20 +150,36 @@ public override string ToString()
 
  public QueryDefinition ToQueryDefinition(IDictionary<object, string> parameters = null)
  {
- SqlQuerySpec querySpec = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions, parameters);
- return QueryDefinition.CreateFromQuerySpec(querySpec);
+ LinqQueryOperation linqQueryOperation = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions, parameters);
+ ScalarOperationKind scalarOperationKind = linqQueryOperation.ScalarOperationKind;
+ Debug.Assert(
+ scalarOperationKind == ScalarOperationKind.None,
+ "CosmosLinqQuery Assert!",
+ $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'");
+
+ return QueryDefinition.CreateFromQuerySpec(linqQueryOperation.SqlQuerySpec);
  }
 
  public FeedIterator<T> ToFeedIterator()
  {
- return new FeedIteratorInlineCore<T>(this.CreateFeedIterator(true),
- this.container.ClientContext);
+ FeedIterator<T> iterator = this.CreateFeedIterator(true, out ScalarOperationKind scalarOperationKind);
+ Debug.Assert(
+ scalarOperationKind == ScalarOperationKind.None,
+ "CosmosLinqQuery Assert!",
+ $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'");
+
+ return new FeedIteratorInlineCore<T>(iterator, this.container.ClientContext);
  }
 
  public FeedIterator ToStreamIterator()
  {
- return new FeedIteratorInlineCore(this.CreateStreamIterator(true),
- this.container.ClientContext);
+ FeedIterator iterator = this.CreateStreamIterator(true, out ScalarOperationKind scalarOperationKind);
+ Debug.Assert(
+ scalarOperationKind == ScalarOperationKind.None,
+ "CosmosLinqQuery Assert!",
+ $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'");
+
+ return new FeedIteratorInlineCore(iterator, this.container.ClientContext);
  }
 
  public void Dispose()
@@ -180,15 +202,18 @@ internal async Task<Response<T>> AggregateResultAsync(CancellationToken cancella
  List<T> result = new List<T>();
  Headers headers = new Headers();
 
- FeedIterator<T> localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false);
- FeedIteratorInternal<T> localFeedIteratorInternal = (FeedIteratorInternal<T>)localFeedIterator;
+ FeedIteratorInlineCore<T> localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false, scalarOperationKind: out ScalarOperationKind scalarOperationKind);
+ Debug.Assert(
+ scalarOperationKind == ScalarOperationKind.None,
+ "CosmosLinqQuery Assert!",
+ $"Unexpected client operation. Expected 'None', Received '{scalarOperationKind}'");
 
  ITrace rootTrace;
  using (rootTrace = Trace.GetRootTrace("Aggregate LINQ Operation"))
  {
  while (localFeedIterator.HasMoreResults)
  {
- FeedResponse<T> response = await localFeedIteratorInternal.ReadNextAsync(rootTrace, cancellationToken);
+ FeedResponse<T> response = await localFeedIterator.ReadNextAsync(rootTrace, cancellationToken);
  headers.RequestCharge += response.RequestCharge;
  result.AddRange(response);
  }
@@ -202,23 +227,57 @@ internal async Task<Response<T>> AggregateResultAsync(CancellationToken cancella
  null);
  }
 
- private FeedIteratorInternal CreateStreamIterator(bool isContinuationExcpected)
+ internal T ExecuteScalar()
+ {
+ FeedIteratorInlineCore<T> localFeedIterator = this.CreateFeedIterator(isContinuationExpected: false, out ScalarOperationKind scalarOperationKind);
+ Headers headers = new Headers();
+
+ List<T> result = new List<T>();
+ ITrace rootTrace;
+ using (rootTrace = Trace.GetRootTrace("Scalar LINQ Operation"))
+ {
+ while (localFeedIterator.HasMoreResults)
+ {
+ FeedResponse<T> response = localFeedIterator.ReadNextAsync(rootTrace, cancellationToken: default).GetAwaiter().GetResult();
+ headers.RequestCharge += response.RequestCharge;
+ result.AddRange(response);
+ }
+ }
+
+ switch (scalarOperationKind)
+ {
+ case ScalarOperationKind.FirstOrDefault:
+ return result.FirstOrDefault();
+
+ // ExecuteScalar gets called when (sync) aggregates such as Max, Min, Sum are invoked on the IQueryable.
+ // Since query fully supprots these operations, there is no client operation involved.
+ // In these cases we return FirstOrDefault which handles empty/undefined/null result set from the backend.
+ case ScalarOperationKind.None:
+ return result.SingleOrDefault();
+
+ default:
+ throw new InvalidOperationException($"Unsupported scalar operation {scalarOperationKind}");
+ }
+ }
+
+ private FeedIteratorInternal CreateStreamIterator(bool isContinuationExcpected, out ScalarOperationKind scalarOperationKind)
  {
- SqlQuerySpec querySpec = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions);
+ LinqQueryOperation linqQueryOperation = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions);
+ scalarOperationKind = linqQueryOperation.ScalarOperationKind;
 
  return this.container.GetItemQueryStreamIteratorInternal(
- sqlQuerySpec: querySpec,
+ sqlQuerySpec: linqQueryOperation.SqlQuerySpec,
  isContinuationExcpected: isContinuationExcpected,
  continuationToken: this.continuationToken,
  feedRange: null,
  requestOptions: this.cosmosQueryRequestOptions);
  }
 
- private FeedIterator<T> CreateFeedIterator(bool isContinuationExpected)
+ private FeedIteratorInlineCore<T> CreateFeedIterator(bool isContinuationExpected, out ScalarOperationKind scalarOperationKind)
  {
- SqlQuerySpec querySpec = DocumentQueryEvaluator.Evaluate(this.Expression, this.linqSerializationOptions);
-
- FeedIteratorInternal streamIterator = this.CreateStreamIterator(isContinuationExpected);
+ FeedIteratorInternal streamIterator = this.CreateStreamIterator(
+ isContinuationExpected,
+  out scalarOperationKind);
  return new FeedIteratorInlineCore<T>(new FeedIteratorCore<T>(
  streamIterator,
  this.responseFactory.CreateQueryFeedUserTypeResponse<T>),

Microsoft.Azure.Cosmos/src/Linq/CosmosLinqQueryProvider.cs

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos.Linq
 {
  using System;
+ using System.Collections.Generic;
  using System.Linq;
  using System.Linq.Expressions;
  using System.Threading;
@@ -60,6 +61,7 @@ public IQueryable<TElement> CreateQuery<TElement>(Expression expression)
 
  public IQueryable CreateQuery(Expression expression)
  {
+ // ISSUE-TODO-adityasa-2024/1/26 - Investigate if reflection usage can be removed.
  Type expressionType = TypeSystem.GetElementType(expression.Type);
  Type documentQueryType = typeof(CosmosLinqQuery<bool>).GetGenericTypeDefinition().MakeGenericType(expressionType);
  return (IQueryable)Activator.CreateInstance(
@@ -76,6 +78,7 @@ public IQueryable CreateQuery(Expression expression)
 
  public TResult Execute<TResult>(Expression expression)
  {
+ // ISSUE-TODO-adityasa-2024/1/26 - We should be able to delegate the implementation to ExecuteAggregateAsync method below by providing an Async implementation of ExecuteScalar.
  Type cosmosQueryType = typeof(CosmosLinqQuery<bool>).GetGenericTypeDefinition().MakeGenericType(typeof(TResult));
  CosmosLinqQuery<TResult> cosmosLINQQuery = (CosmosLinqQuery<TResult>)Activator.CreateInstance(
  cosmosQueryType,
@@ -88,7 +91,7 @@ public TResult Execute<TResult>(Expression expression)
  this.allowSynchronousQueryExecution,
  this.linqSerializerOptions);
  this.onExecuteScalarQueryCallback?.Invoke(cosmosLINQQuery);
- return cosmosLINQQuery.ToList().FirstOrDefault();
+ return cosmosLINQQuery.ExecuteScalar();
  }
 
  //Sync execution of query via direct invoke on IQueryProvider.

Microsoft.Azure.Cosmos/src/Linq/DocumentQuery.cs

@@ -273,10 +273,10 @@ IEnumerator IEnumerable.GetEnumerator()
 
  public override string ToString()
  {
- SqlQuerySpec querySpec = DocumentQueryEvaluator.Evaluate(this.Expression);
- if (querySpec != null)
+ LinqQueryOperation linqQueryOperation = DocumentQueryEvaluator.Evaluate(this.Expression);
+ if (linqQueryOperation.SqlQuerySpec != null)
  {
- return JsonConvert.SerializeObject(querySpec);
+ return JsonConvert.SerializeObject(linqQueryOperation.SqlQuerySpec);
  }
 
  return new Uri(this.client.ServiceEndpoint, this.documentsFeedOrDatabaseLink).ToString();

Microsoft.Azure.Cosmos/src/Linq/DocumentQueryEvaluator.cs

@@ -15,7 +15,7 @@ internal static class DocumentQueryEvaluator
  {
  private const string SQLMethod = "AsSQL";
 
- public static SqlQuerySpec Evaluate(
+ public static LinqQueryOperation Evaluate(
  Expression expression,
  CosmosLinqSerializerOptionsInternal linqSerializerOptions = null,
  IDictionary<object, string> parameters = null)
@@ -51,7 +51,7 @@ public static bool IsTransformExpression(Expression expression)
  /// foreach(Database db in client.CreateDatabaseQuery()) {} 
  /// </summary>
  /// <param name="expression"></param>
- private static SqlQuerySpec HandleEmptyQuery(ConstantExpression expression)
+ private static LinqQueryOperation HandleEmptyQuery(ConstantExpression expression)
  {
  if (expression.Value == null)
  {
@@ -69,11 +69,12 @@ private static SqlQuerySpec HandleEmptyQuery(ConstantExpression expression)
  ClientResources.BadQuery_InvalidExpression,
  expression.ToString()));
  }
+
  //No query specified.
- return null;
+ return new LinqQueryOperation(sqlQuerySpec: null, scalarOperationKind: ScalarOperationKind.None);
  }
 
- private static SqlQuerySpec HandleMethodCallExpression(
+ private static LinqQueryOperation HandleMethodCallExpression(
  MethodCallExpression expression,
  IDictionary<object, string> parameters,
  CosmosLinqSerializerOptionsInternal linqSerializerOptions = null)
@@ -100,7 +101,7 @@ private static SqlQuerySpec HandleMethodCallExpression(
  /// foreach(string record in client.CreateDocumentQuery().Navigate("Raw JQuery"))
  /// </summary>
  /// <param name="expression"></param>
- private static SqlQuerySpec HandleAsSqlTransformExpression(MethodCallExpression expression)
+ private static LinqQueryOperation HandleAsSqlTransformExpression(MethodCallExpression expression)
  {
  Expression paramExpression = expression.Arguments[1];
 
@@ -122,7 +123,7 @@ private static SqlQuerySpec HandleAsSqlTransformExpression(MethodCallExpression
  }
  }
 
- private static SqlQuerySpec GetSqlQuerySpec(object value)
+ private static LinqQueryOperation GetSqlQuerySpec(object value)
  {
  if (value == null)
  {
@@ -133,11 +134,11 @@ private static SqlQuerySpec GetSqlQuerySpec(object value)
  }
  else if (value.GetType() == typeof(SqlQuerySpec))
  {
- return (SqlQuerySpec)value;
+ return new LinqQueryOperation((SqlQuerySpec)value, ScalarOperationKind.None);
  }
  else if (value.GetType() == typeof(string))
  {
- return new SqlQuerySpec((string)value);
+ return new LinqQueryOperation(new SqlQuerySpec((string)value), ScalarOperationKind.None);
  }
  else
  {