Query: Introduce AsSplitQuery operator to run collection projection i…

…n separate command for relational

Resolves #20892

src/EFCore.Cosmos/Query/Internal/CosmosQueryMetadataExtractingExpressionVisitor.cs

@@ -49,6 +49,5 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
 
             return base.VisitMethodCall(methodCallExpression);
         }
-
     }
 }

src/EFCore.Cosmos/Query/Internal/CosmosQueryTranslationPreprocessor.cs

@@ -39,9 +39,8 @@ public CosmosQueryTranslationPreprocessor(
         /// </summary>
         public override Expression NormalizeQueryableMethod(Expression query)
         {
-            query = base.NormalizeQueryableMethod(query);
-
             query = new CosmosQueryMetadataExtractingExpressionVisitor(_queryCompilationContext).Visit(query);
+            query = base.NormalizeQueryableMethod(query);
 
             return query;
         }

src/EFCore.Relational/Extensions/RelationalQueryableExtensions.cs

@@ -6,6 +6,7 @@
 using System.Data.Common;
 using System.Linq;
 using System.Linq.Expressions;
+using System.Reflection;
 using JetBrains.Annotations;
 using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Query;
@@ -92,7 +93,7 @@ public static DbCommand CreateDbCommand([NotNull] this IQueryable source)
         [StringFormatMethod("sql")]
         public static IQueryable<TEntity> FromSqlRaw<TEntity>(
             [NotNull] this DbSet<TEntity> source,
-            [NotNull] [NotParameterized] string sql,
+            [NotNull][NotParameterized] string sql,
             [NotNull] params object[] parameters)
             where TEntity : class
         {
@@ -133,7 +134,7 @@ public static IQueryable<TEntity> FromSqlRaw<TEntity>(
         /// <returns> An <see cref="IQueryable{T}" /> representing the interpolated string SQL query. </returns>
         public static IQueryable<TEntity> FromSqlInterpolated<TEntity>(
             [NotNull] this DbSet<TEntity> source,
-            [NotNull] [NotParameterized] FormattableString sql)
+            [NotNull][NotParameterized] FormattableString sql)
             where TEntity : class
         {
             Check.NotNull(source, nameof(source));
@@ -161,5 +162,31 @@ private static FromSqlQueryRootExpression GenerateFromSqlQueryRoot(
                 sql,
                 Expression.Constant(arguments));
         }
+
+        /// <summary>
+        ///     <para>
+        ///         Returns a new query in which the collections in the query results will be loaded through separate database queries.
+        ///     </para>
+        ///     <para>
+        ///         This strategy fetches all the data from the server through separate database queries before generating any results.
+        ///     </para>
+        /// </summary>
+        /// <typeparam name="TEntity"> The type of entity being queried. </typeparam>
+        /// <param name="source"> The source query. </param>
+        /// <returns> A new query where collections will be loaded through separate database queries. </returns>
+        public static IQueryable<TEntity> AsSplitQuery<TEntity>(
+            [NotNull] this IQueryable<TEntity> source)
+            where TEntity : class
+        {
+            Check.NotNull(source, nameof(source));
+
+            return source.Provider is EntityQueryProvider
+                ? source.Provider.CreateQuery<TEntity>(
+                    Expression.Call(AsSplitQueryMethodInfo.MakeGenericMethod(typeof(TEntity)), source.Expression))
+                : source;
+        }
+
+        internal static readonly MethodInfo AsSplitQueryMethodInfo
+            = typeof(RelationalQueryableExtensions).GetTypeInfo().GetDeclaredMethod(nameof(AsSplitQuery));
     }
 }

src/EFCore.Relational/Infrastructure/EntityFrameworkRelationalServicesBuilder.cs

@@ -169,6 +169,7 @@ public override EntityFrameworkServicesBuilder TryAddCoreServices()
             TryAdd<IQueryTranslationPreprocessorFactory, RelationalQueryTranslationPreprocessorFactory>();
             TryAdd<IRelationalParameterBasedSqlProcessorFactory, RelationalParameterBasedSqlProcessorFactory>();
             TryAdd<IRelationalQueryStringFactory, RelationalQueryStringFactory>();
+            TryAdd<IQueryCompilationContextFactory, RelationalQueryCompilationContextFactory>();
 
             ServiceCollectionMap.GetInfrastructure()
                 .AddDependencySingleton<RelationalSqlGenerationHelperDependencies>()
@@ -201,7 +202,8 @@ public override EntityFrameworkServicesBuilder TryAddCoreServices()
                 .AddDependencyScoped<RelationalCompiledQueryCacheKeyGeneratorDependencies>()
                 .AddDependencyScoped<RelationalConnectionDependencies>()
                 .AddDependencyScoped<RelationalDatabaseDependencies>()
-                .AddDependencyScoped<RelationalQueryContextDependencies>();
+                .AddDependencyScoped<RelationalQueryContextDependencies>()
+                .AddDependencyScoped<RelationalQueryCompilationContextDependencies>();
 
             return base.TryAddCoreServices();
         }

src/EFCore.Relational/Query/Internal/CollectionJoinApplyingExpressionVisitor.cs

@@ -15,8 +15,20 @@ namespace Microsoft.EntityFrameworkCore.Query.Internal
     /// </summary>
     public class CollectionJoinApplyingExpressionVisitor : ExpressionVisitor
     {
+        private readonly bool _splitQuery;
         private int _collectionId;
 
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        public CollectionJoinApplyingExpressionVisitor(bool splitQuery)
+        {
+            _splitQuery = splitQuery;
+        }
+
         /// <summary>
         ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
         ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -48,7 +60,8 @@ protected override Expression VisitExtension(Expression extensionExpression)
                     collectionId,
                     innerShaper,
                     collectionShaperExpression.Navigation,
-                    collectionShaperExpression.ElementType);
+                    collectionShaperExpression.ElementType,
+                    _splitQuery);
             }
 
             return extensionExpression is ShapedQueryExpression shapedQueryExpression

src/EFCore.Relational/Query/Internal/QueryingEnumerable.cs

@@ -23,7 +23,7 @@ public class QueryingEnumerable<T> : IEnumerable<T>, IAsyncEnumerable<T>, IRelat
     {
         private readonly RelationalQueryContext _relationalQueryContext;
         private readonly RelationalCommandCache _relationalCommandCache;
-        private readonly Func<QueryContext, DbDataReader, ResultContext, ResultCoordinator, T> _shaper;
+        private readonly Func<QueryContext, DbDataReader, ResultContext, SingleQueryResultCoordinator, T> _shaper;
         private readonly Type _contextType;
         private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
         private readonly bool _performIdentityResolution;
@@ -37,7 +37,7 @@ public class QueryingEnumerable<T> : IEnumerable<T>, IAsyncEnumerable<T>, IRelat
         public QueryingEnumerable(
             [NotNull] RelationalQueryContext relationalQueryContext,
             [NotNull] RelationalCommandCache relationalCommandCache,
-            [NotNull] Func<QueryContext, DbDataReader, ResultContext, ResultCoordinator, T> shaper,
+            [NotNull] Func<QueryContext, DbDataReader, ResultContext, SingleQueryResultCoordinator, T> shaper,
             [NotNull] Type contextType,
             bool performIdentityResolution)
         {
@@ -106,13 +106,13 @@ private sealed class Enumerator : IEnumerator<T>
         {
             private readonly RelationalQueryContext _relationalQueryContext;
             private readonly RelationalCommandCache _relationalCommandCache;
-            private readonly Func<QueryContext, DbDataReader, ResultContext, ResultCoordinator, T> _shaper;
+            private readonly Func<QueryContext, DbDataReader, ResultContext, SingleQueryResultCoordinator, T> _shaper;
             private readonly Type _contextType;
             private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
             private readonly bool _performIdentityResolution;
 
             private RelationalDataReader _dataReader;
-            private ResultCoordinator _resultCoordinator;
+            private SingleQueryResultCoordinator _resultCoordinator;
             private IExecutionStrategy _executionStrategy;
 
             public Enumerator(QueryingEnumerable<T> queryingEnumerable)
@@ -200,7 +200,7 @@ private bool InitializeReader(DbContext _, bool result)
                             _relationalQueryContext.Context,
                             _relationalQueryContext.CommandLogger));
 
-                _resultCoordinator = new ResultCoordinator();
+                _resultCoordinator = new SingleQueryResultCoordinator();
 
                 _relationalQueryContext.InitializeStateManager(_performIdentityResolution);
 
@@ -220,14 +220,14 @@ private sealed class AsyncEnumerator : IAsyncEnumerator<T>
         {
             private readonly RelationalQueryContext _relationalQueryContext;
             private readonly RelationalCommandCache _relationalCommandCache;
-            private readonly Func<QueryContext, DbDataReader, ResultContext, ResultCoordinator, T> _shaper;
+            private readonly Func<QueryContext, DbDataReader, ResultContext, SingleQueryResultCoordinator, T> _shaper;
             private readonly Type _contextType;
             private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
             private readonly bool _performIdentityResolution;
             private readonly CancellationToken _cancellationToken;
 
             private RelationalDataReader _dataReader;
-            private ResultCoordinator _resultCoordinator;
+            private SingleQueryResultCoordinator _resultCoordinator;
             private IExecutionStrategy _executionStrategy;
 
             public AsyncEnumerator(
@@ -319,7 +319,7 @@ private async Task<bool> InitializeReaderAsync(DbContext _, bool result, Cancell
                         cancellationToken)
                         .ConfigureAwait(false);
 
-                _resultCoordinator = new ResultCoordinator();
+                _resultCoordinator = new SingleQueryResultCoordinator();
 
                 _relationalQueryContext.InitializeStateManager(_performIdentityResolution);
 

src/EFCore.Relational/Query/Internal/RelationalQueryCompilationContextFactory.cs

@@ -0,0 +1,56 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Utilities;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.EntityFrameworkCore.Query.Internal
+{
+    /// <summary>
+    ///     <para>
+    ///         This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///         the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///         any release. You should only use it directly in your code with extreme caution and knowing that
+    ///         doing so can result in application failures when updating to a new Entity Framework Core release.
+    ///     </para>
+    ///     <para>
+    ///         The service lifetime is <see cref="ServiceLifetime.Scoped" />. This means that each
+    ///         <see cref="DbContext" /> instance will use its own instance of this service.
+    ///         The implementation may depend on other services registered with any lifetime.
+    ///         The implementation does not need to be thread-safe.
+    ///     </para>
+    /// </summary>
+    public class RelationalQueryCompilationContextFactory : IQueryCompilationContextFactory
+    {
+        private readonly QueryCompilationContextDependencies _dependencies;
+        private readonly RelationalQueryCompilationContextDependencies _relationalDependencies;
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        public RelationalQueryCompilationContextFactory(
+            [NotNull] QueryCompilationContextDependencies dependencies,
+            [NotNull] RelationalQueryCompilationContextDependencies relationalDependencies)
+        {
+            Check.NotNull(dependencies, nameof(dependencies));
+            Check.NotNull(relationalDependencies, nameof(relationalDependencies));
+
+            _dependencies = dependencies;
+            _relationalDependencies = relationalDependencies;
+        }
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        public virtual QueryCompilationContext Create(bool async)
+            => new RelationalQueryCompilationContext(_dependencies, _relationalDependencies, async);
+    }
+}
+

src/EFCore.Relational/Query/Internal/RelationalQueryMetadataExtractingExpressionVisitor.cs

@@ -0,0 +1,54 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Linq.Expressions;
+using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Utilities;
+
+namespace Microsoft.EntityFrameworkCore.Query.Internal
+{
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     any release. You should only use it directly in your code with extreme caution and knowing that
+    ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+    /// </summary>
+    public class RelationalQueryMetadataExtractingExpressionVisitor : ExpressionVisitor
+    {
+        private readonly RelationalQueryCompilationContext _relationalQueryCompilationContext;
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        public RelationalQueryMetadataExtractingExpressionVisitor([NotNull] RelationalQueryCompilationContext relationalQueryCompilationContext)
+        {
+            Check.NotNull(relationalQueryCompilationContext, nameof(relationalQueryCompilationContext));
+
+            _relationalQueryCompilationContext = relationalQueryCompilationContext;
+        }
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     any release. You should only use it directly in your code with extreme caution and knowing that
+        ///     doing so can result in application failures when updating to a new Entity Framework Core release.
+        /// </summary>
+        protected override Expression VisitMethodCall(MethodCallExpression methodCallExpression)
+        {
+            if (methodCallExpression.Method.IsGenericMethod
+                && methodCallExpression.Method.GetGenericMethodDefinition() == RelationalQueryableExtensions.AsSplitQueryMethodInfo)
+            {
+                var innerQueryable = Visit(methodCallExpression.Arguments[0]);
+
+                _relationalQueryCompilationContext.IsSplitQuery = true;
+
+                return innerQueryable;
+            }
+
+            return base.VisitMethodCall(methodCallExpression);
+        }
+    }
+}

src/EFCore.Relational/Query/Internal/CollectionMaterializationContext.cs → src/EFCore.Relational/Query/Internal/SingleQueryCollectionContext.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using JetBrains.Annotations;
@@ -11,16 +11,16 @@ namespace Microsoft.EntityFrameworkCore.Query.Internal
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public class CollectionMaterializationContext
+    public class SingleQueryCollectionContext
     {
         /// <summary>
         ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
         ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
         ///     any release. You should only use it directly in your code with extreme caution and knowing that
         ///     doing so can result in application failures when updating to a new Entity Framework Core release.
         /// </summary>
-        public CollectionMaterializationContext(
-            [NotNull] object parent,
+        public SingleQueryCollectionContext(
+            [CanBeNull] object parent,
             [NotNull] object collection,
             [NotNull] object[] parentIdentifier,
             [NotNull] object[] outerIdentifier)