Query: Avoid pulling un-necessary columns in TPH after OfType

src/EFCore.Cosmos/Query/Expressions/Internal/SelectExpression.cs

@@ -40,7 +40,7 @@ public BinaryExpression GetDiscriminatorPredicate(IEntityType entityType)
             }
 
             var concreteEntityTypes
-                = entityType.GetConcreteTypesInHierarchy().ToList();
+                = entityType.GetConcreteDerivedTypesInclusive().ToList();
 
             var discriminatorProperty = entityType.GetDiscriminatorProperty();
 

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

@@ -101,7 +101,7 @@ private LambdaExpression CreateMaterializerExpression(
             var materializationContextParameter
                 = Expression.Parameter(typeof(MaterializationContext), "materializationContext");
 
-            var concreteEntityTypes = entityType.GetConcreteTypesInHierarchy().ToList();
+            var concreteEntityTypes = entityType.GetConcreteDerivedTypesInclusive().ToList();
             var firstEntityType = concreteEntityTypes[0];
             var indexMap = new int[firstEntityType.PropertyCount()];
 

src/EFCore.Relational/Query/Pipeline/EntityProjectionExpression.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Linq.Expressions;
 using Microsoft.EntityFrameworkCore.Metadata;
+using Microsoft.EntityFrameworkCore.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Relational.Query.Pipeline.SqlExpressions;
 
 namespace Microsoft.EntityFrameworkCore.Relational.Query.Pipeline
@@ -18,14 +20,14 @@ private readonly IDictionary<IProperty, ColumnExpression> _propertyExpressionsCa
 
         public EntityProjectionExpression(IEntityType entityType, TableExpressionBase innerTable, bool nullable)
         {
-            EntityType = entityType.RootType();
+            EntityType = entityType;
             _innerTable = innerTable;
             _nullable = nullable;
         }
 
         public EntityProjectionExpression(IEntityType entityType, IDictionary<IProperty, ColumnExpression> propertyExpressions)
         {
-            EntityType = entityType.RootType();
+            EntityType = entityType;
             _propertyExpressionsCache = propertyExpressions;
         }
 
@@ -75,13 +77,23 @@ public EntityProjectionExpression MakeNullable()
             }
         }
 
+        public EntityProjectionExpression UpdateEntityType(IEntityType derivedType)
+        {
+            if (_innerTable != null)
+            {
+                return new EntityProjectionExpression(derivedType, _innerTable, _nullable);
+            }
+
+            throw new InvalidOperationException("EntityProjectionExpression: Cannot update EntityType when _innerTable is null");
+        }
+
         public IEntityType EntityType { get; }
         public override ExpressionType NodeType => ExpressionType.Extension;
         public override Type Type => EntityType.ClrType;
 
         public ColumnExpression GetProperty(IProperty property)
         {
-            if (property.DeclaringEntityType.RootType() != EntityType)
+            if (!EntityType.GetTypesInHierarchy().Contains(property.DeclaringEntityType))
             {
                 throw new InvalidOperationException("Called EntityProjectionExpression.GetProperty() with incorrect IProperty");
             }

src/EFCore.Relational/Query/Pipeline/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -14,6 +14,7 @@
 using Microsoft.EntityFrameworkCore.Query.Pipeline;
 using Microsoft.EntityFrameworkCore.Relational.Query.Pipeline.SqlExpressions;
 using System.IO;
+using System.Diagnostics;
 
 namespace Microsoft.EntityFrameworkCore.Relational.Query.Pipeline
 {
@@ -559,7 +560,7 @@ protected override ShapedQueryExpression TranslateOfType(ShapedQueryExpression s
                 if (derivedType != null)
                 {
                     var selectExpression = (SelectExpression)source.QueryExpression;
-                    var concreteEntityTypes = derivedType.GetConcreteTypesInHierarchy().ToList();
+                    var concreteEntityTypes = derivedType.GetConcreteDerivedTypesInclusive().ToList();
                     var discriminatorColumn = selectExpression
                         .BindProperty(entityShaperExpression.ValueBufferExpression, entityType.GetDiscriminatorProperty());
 
@@ -572,6 +573,19 @@ protected override ShapedQueryExpression TranslateOfType(ShapedQueryExpression s
 
                     selectExpression.ApplyPredicate(predicate);
 
+                    var projectionMember = entityShaperExpression.ValueBufferExpression.ProjectionMember;
+
+                    Debug.Assert(new ProjectionMember().Equals(projectionMember),
+                        "Invalid ProjectionMember when processing OfType");
+
+                    var entityProjection = (EntityProjectionExpression)selectExpression.GetMappedProjection(projectionMember);
+
+                    selectExpression.ReplaceProjectionMapping(
+                        new Dictionary<ProjectionMember, Expression>
+                        {
+                            { projectionMember, entityProjection.UpdateEntityType(derivedType)}
+                        });
+
                     source.ShaperExpression = new EntityShaperExpression(
                         derivedType, entityShaperExpression.ValueBufferExpression, entityShaperExpression.Nullable);
 

src/EFCore.Relational/Query/Pipeline/RelationalShapedQueryCompilingExpressionVisitor.cs

@@ -21,7 +21,7 @@ public partial class RelationalShapedQueryCompilingExpressionVisitor : ShapedQue
         private readonly IParameterNameGeneratorFactory _parameterNameGeneratorFactory;
         private readonly Type _contextType;
         private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _logger;
-        private static ParameterExpression _resultCoordinatorParameter
+        private static readonly ParameterExpression _resultCoordinatorParameter
             = Expression.Parameter(typeof(ResultCoordinator), "resultCoordinator");
 
         public RelationalShapedQueryCompilingExpressionVisitor(
@@ -111,7 +111,7 @@ protected override Expression VisitShapedQueryExpression(ShapedQueryExpression s
 
         private class IndexMapInjectingExpressionVisitor : ExpressionVisitor
         {
-            private ParameterExpression _indexMapParameter;
+            private readonly ParameterExpression _indexMapParameter;
 
             public IndexMapInjectingExpressionVisitor(ParameterExpression indexMapParameter)
             {

src/EFCore.Relational/Query/Pipeline/RelationalSqlTranslatingExpressionVisitor.cs

@@ -145,7 +145,7 @@ protected override Expression VisitTypeBinary(TypeBinaryExpression typeBinaryExp
                 var derivedType = entityType.GetDerivedTypes().SingleOrDefault(et => et.ClrType == typeBinaryExpression.TypeOperand);
                 if (derivedType != null)
                 {
-                    var concreteEntityTypes = derivedType.GetConcreteTypesInHierarchy().ToList();
+                    var concreteEntityTypes = derivedType.GetConcreteDerivedTypesInclusive().ToList();
                     var discriminatorColumn = BindProperty(entityShaperExpression, entityType.GetDiscriminatorProperty());
 
                     return concreteEntityTypes.Count == 1

src/EFCore.Relational/Query/Pipeline/SqlExpressionFactory.cs

@@ -539,7 +539,7 @@ public SelectExpression Select(IEntityType entityType, string sql, Expression sq
 
         private void AddDiscriminator(SelectExpression selectExpression, IEntityType entityType)
         {
-            var concreteEntityTypes = entityType.GetConcreteTypesInHierarchy().ToList();
+            var concreteEntityTypes = entityType.GetConcreteDerivedTypesInclusive().ToList();
 
             if (concreteEntityTypes.Count == 1)
             {

src/EFCore.Relational/Query/Pipeline/SqlExpressions/SelectExpression.cs

@@ -8,6 +8,7 @@
 using System.Linq.Expressions;
 using System.Reflection;
 using Microsoft.EntityFrameworkCore.Metadata;
+using Microsoft.EntityFrameworkCore.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Query.Internal;
 using Microsoft.EntityFrameworkCore.Query.Pipeline;
 
@@ -123,8 +124,8 @@ public void ApplyProjection()
                 if (keyValuePair.Value is EntityProjectionExpression entityProjection)
                 {
                     var map = new Dictionary<IProperty, int>();
-                    foreach (var property in entityProjection.EntityType
-                        .GetDerivedTypesInclusive().SelectMany(e => e.GetDeclaredProperties()))
+
+                    foreach (var property in GetAllPropertiesInHierarchy(entityProjection.EntityType))
                     {
                         map[property] = AddToProjection(entityProjection.GetProperty(property));
                     }
@@ -140,6 +141,11 @@ public void ApplyProjection()
             _projectionMapping = result;
         }
 
+        private IEnumerable<IProperty> GetAllPropertiesInHierarchy(IEntityType entityType)
+        {
+            return entityType.GetTypesInHierarchy().SelectMany(e => e.GetDeclaredProperties());
+        }
+
         public void ReplaceProjectionMapping(IDictionary<ProjectionMember, Expression> projectionMapping)
         {
             _projectionMapping.Clear();
@@ -188,8 +194,7 @@ public IDictionary<IProperty, int> AddToProjection(EntityProjectionExpression en
             if (!_entityProjectionCache.TryGetValue(entityProjection, out var dictionary))
             {
                 dictionary = new Dictionary<IProperty, int>();
-                foreach (var property in entityProjection.EntityType
-                    .GetDerivedTypesInclusive().SelectMany(e => e.GetDeclaredProperties()))
+                foreach (var property in GetAllPropertiesInHierarchy(entityProjection.EntityType))
                 {
                     dictionary[property] = AddToProjection(entityProjection.GetProperty(property));
                 }
@@ -293,23 +298,6 @@ public void ClearOrdering()
             _orderings.Clear();
         }
 
-        private string GenerateUniqueName(HashSet<string> usedNames, string prefix)
-        {
-            if (!usedNames.Contains(prefix))
-            {
-                return prefix;
-            }
-
-            var counter = 0;
-            var uniqueName = prefix + counter;
-            while (usedNames.Contains(uniqueName))
-            {
-                uniqueName = prefix + counter++;
-            }
-
-            return uniqueName;
-        }
-
         public SelectExpression PushdownIntoSubQuery()
         {
             var subquery = new SelectExpression("t", new List<ProjectionExpression>(), _tables.ToList(), _orderings.ToList())
@@ -346,8 +334,7 @@ public SelectExpression PushdownIntoSubQuery()
                     if (mapping.Value is EntityProjectionExpression entityProjection)
                     {
                         var propertyExpressions = new Dictionary<IProperty, ColumnExpression>();
-                        foreach (var property in entityProjection.EntityType
-                            .GetDerivedTypesInclusive().SelectMany(e => e.GetDeclaredProperties()))
+                        foreach (var property in GetAllPropertiesInHierarchy(entityProjection.EntityType))
                         {
                             var innerColumn = entityProjection.GetProperty(property);
                             var index = subquery.AddToProjection(innerColumn);

src/EFCore/Metadata/Internal/EntityTypeExtensions.cs

@@ -384,9 +384,18 @@ public static Func<ISnapshot> GetEmptyShadowValuesFactory([NotNull] this IEntity
         ///     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 static IEnumerable<IEntityType> GetConcreteTypesInHierarchy([NotNull] this IEntityType entityType)
+        public static IEnumerable<IEntityType> GetConcreteDerivedTypesInclusive([NotNull] this IEntityType entityType)
             => entityType.GetDerivedTypesInclusive().Where(et => !et.IsAbstract());
 
+        /// <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 static IEnumerable<IEntityType> GetTypesInHierarchy([NotNull] this IEntityType entityType)
+            => entityType.GetAllBaseTypes().Concat(entityType.GetDerivedTypesInclusive());
+
         /// <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

src/EFCore/Query/Pipeline/QueryableMethodTranslatingExpressionVisitor.cs

@@ -493,12 +493,14 @@ protected ShapedQueryExpression TranslateResultSelectorForJoin(
         }
 
         protected ShapedQueryExpression TranslateResultSelectorForGroupJoin(
+#pragma warning disable IDE0060 // Remove unused parameter
             ShapedQueryExpression outer,
             Expression innerShaper,
             LambdaExpression outerKeySelector,
             LambdaExpression innerKeySelector,
             LambdaExpression resultSelector,
             Type transparentIdentifierType)
+#pragma warning restore IDE0060 // Remove unused parameter
         {
             throw new NotImplementedException();
         }

src/EFCore/Query/Pipeline/ShapedQueryExpressionVisitor.cs

@@ -327,7 +327,7 @@ var valueBufferConstructor
                                 QueryCompilationContext.QueryContextParameter,
                                 _dbContextMemberInfo))));
 
-                var concreteEntityTypes = entityType.GetConcreteTypesInHierarchy().ToList();
+                var concreteEntityTypes = entityType.GetConcreteDerivedTypesInclusive().ToList();
 
                 Expression materializationExpression = null;
                 Expression shadowValuesExpression = Expression.Constant(ValueBuffer.Empty);