Query: Assign nullability to entity shaper correctly for set operatio…

…ns (#23227)

Also improved In-Memory set operation implementation

Resolves #19253

src/EFCore.InMemory/Query/Internal/InMemoryQueryExpression.cs

@@ -379,6 +379,110 @@ public virtual void PushdownIntoSubquery()
             }
         }
 
+        /// <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 void ApplySetOperation([NotNull] MethodInfo setOperationMethodInfo, [NotNull] InMemoryQueryExpression source2)
+        {
+            var clientProjection = _valueBufferSlots.Count != 0;
+            if (!clientProjection)
+            {
+                var result = new Dictionary<ProjectionMember, Expression>();
+                foreach (var (key, value1, value2) in _projectionMapping.Join(
+                    source2._projectionMapping, kv => kv.Key, kv => kv.Key,
+                    (kv1, kv2) => (kv1.Key, Value1: kv1.Value, Value2: kv2.Value)))
+                {
+                    if (value1 is EntityProjectionExpression entityProjection1
+                        && value2 is EntityProjectionExpression entityProjection2)
+                    {
+                        var map = new Dictionary<IProperty, Expression>();
+                        foreach (var property in GetAllPropertiesInHierarchy(entityProjection1.EntityType))
+                        {
+                            var expressionToAdd1 = entityProjection1.BindProperty(property);
+                            var expressionToAdd2 = entityProjection2.BindProperty(property);
+                            var index = AddToProjection(expressionToAdd1);
+                            source2.AddToProjection(expressionToAdd2);
+                            var type = expressionToAdd1.Type;
+                            if (!type.IsNullableType()
+                                && expressionToAdd2.Type.IsNullableType())
+                            {
+                                type = expressionToAdd2.Type;
+                            }
+                            map[property] = CreateReadValueExpression(type, index, property);
+                        }
+
+                        result[key] = new EntityProjectionExpression(entityProjection1.EntityType, map);
+                    }
+                    else
+                    {
+                        var index = AddToProjection(value1);
+                        source2.AddToProjection(value2);
+                        var type = value1.Type;
+                        if (!type.IsNullableType()
+                            && value2.Type.IsNullableType())
+                        {
+                            type = value2.Type;
+                        }
+                        result[key] = CreateReadValueExpression(type, index, InferPropertyFromInner(value1));
+                    }
+                }
+
+                _projectionMapping = result;
+            }
+
+            var selectorLambda = Lambda(
+                New(
+                    _valueBufferConstructor,
+                    NewArrayInit(
+                        typeof(object),
+                        _valueBufferSlots
+                            .Select(e => e.Type.IsValueType ? Convert(e, typeof(object)) : e))),
+                CurrentParameter);
+
+            _groupingParameter = null;
+
+            ServerQueryExpression = Call(
+                EnumerableMethods.Select.MakeGenericMethod(ServerQueryExpression.Type.TryGetSequenceType(), typeof(ValueBuffer)),
+                ServerQueryExpression,
+                selectorLambda);
+
+            var selectorLambda2 = Lambda(
+                New(
+                    _valueBufferConstructor,
+                    NewArrayInit(
+                        typeof(object),
+                        source2._valueBufferSlots
+                            .Select(e => e.Type.IsValueType ? Convert(e, typeof(object)) : e))),
+                source2.CurrentParameter);
+
+            source2._groupingParameter = null;
+
+            source2.ServerQueryExpression = Call(
+                EnumerableMethods.Select.MakeGenericMethod(source2.ServerQueryExpression.Type.TryGetSequenceType(), typeof(ValueBuffer)),
+                source2.ServerQueryExpression,
+                selectorLambda2);
+
+            ServerQueryExpression = Call(
+                setOperationMethodInfo.MakeGenericMethod(typeof(ValueBuffer)), ServerQueryExpression, source2.ServerQueryExpression);
+
+            if (clientProjection)
+            {
+                var newValueBufferSlots = _valueBufferSlots
+                    .Select((e, i) => CreateReadValueExpression(e.Type, i, InferPropertyFromInner(e)))
+                    .ToList();
+
+                _valueBufferSlots.Clear();
+                _valueBufferSlots.AddRange(newValueBufferSlots);
+            }
+            else
+            {
+                _valueBufferSlots.Clear();
+            }
+        }
+
         /// <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.InMemory/Query/Internal/InMemoryQueryableMethodTranslatingExpressionVisitor.cs

@@ -1588,18 +1588,61 @@ private ShapedQueryExpression TranslateSetOperation(
             var inMemoryQueryExpression1 = (InMemoryQueryExpression)source1.QueryExpression;
             var inMemoryQueryExpression2 = (InMemoryQueryExpression)source2.QueryExpression;
 
-            // Apply any pending selectors, ensuring that the shape of both expressions is identical
-            // prior to applying the set operation.
-            inMemoryQueryExpression1.PushdownIntoSubquery();
-            inMemoryQueryExpression2.PushdownIntoSubquery();
+            inMemoryQueryExpression1.ApplySetOperation(setOperationMethodInfo, inMemoryQueryExpression2);
+
+            if (setOperationMethodInfo.Equals(EnumerableMethods.Except))
+            {
+                return source1;
+            }
+
+            var makeNullable = setOperationMethodInfo != EnumerableMethods.Intersect;
+
+            return source1.UpdateShaperExpression(MatchShaperNullabilityForSetOperation(
+                source1.ShaperExpression, source2.ShaperExpression, makeNullable));
+        }
+
+        private Expression MatchShaperNullabilityForSetOperation(Expression shaper1, Expression shaper2, bool makeNullable)
+        {
+            switch (shaper1)
+            {
+                case EntityShaperExpression entityShaperExpression1
+                when shaper2 is EntityShaperExpression entityShaperExpression2:
+                    return entityShaperExpression1.IsNullable != entityShaperExpression2.IsNullable
+                        ? entityShaperExpression1.MakeNullable(makeNullable)
+                        : entityShaperExpression1;
+
+                case NewExpression newExpression1
+                when shaper2 is NewExpression newExpression2:
+                    var newArguments = new Expression[newExpression1.Arguments.Count];
+                    for (var i = 0; i < newArguments.Length; i++)
+                    {
+                        newArguments[i] = MatchShaperNullabilityForSetOperation(
+                            newExpression1.Arguments[i], newExpression2.Arguments[i], makeNullable);
+                    }
+
+                    return newExpression1.Update(newArguments);
+
+                case MemberInitExpression memberInitExpression1
+                when shaper2 is MemberInitExpression memberInitExpression2:
+                    var newExpression = (NewExpression)MatchShaperNullabilityForSetOperation(
+                        memberInitExpression1.NewExpression, memberInitExpression2.NewExpression, makeNullable);
+
+                    var memberBindings = new MemberBinding[memberInitExpression1.Bindings.Count];
+                    for (var i = 0; i < memberBindings.Length; i++)
+                    {
+                        var memberAssignment = memberInitExpression1.Bindings[i] as MemberAssignment;
+                        Check.DebugAssert(memberAssignment != null, "Only member assignment bindings are supported");
 
-            inMemoryQueryExpression1.UpdateServerQueryExpression(
-                Expression.Call(
-                    setOperationMethodInfo.MakeGenericMethod(typeof(ValueBuffer)),
-                    inMemoryQueryExpression1.ServerQueryExpression,
-                    inMemoryQueryExpression2.ServerQueryExpression));
 
-            return source1;
+                        memberBindings[i] = memberAssignment.Update(MatchShaperNullabilityForSetOperation(
+                            memberAssignment.Expression, ((MemberAssignment)memberInitExpression2.Bindings[i]).Expression, makeNullable));
+                    }
+
+                    return memberInitExpression1.Update(newExpression, memberBindings);
+
+                default:
+                    return shaper1;
+            }
         }
     }
 }

src/EFCore.Relational/Query/RelationalEntityShaperExpression.cs

@@ -160,11 +160,18 @@ public override EntityShaperExpression WithEntityType(IEntityType entityType)
         }
 
         /// <inheritdoc />
+        [Obsolete("Use MakeNullable() instead.")]
         public override EntityShaperExpression MarkAsNullable()
-            => !IsNullable
+            => MakeNullable();
+
+        /// <inheritdoc />
+        public override EntityShaperExpression MakeNullable(bool nullable = true)
+        {
+            return IsNullable != nullable
                 // Marking nullable requires recomputation of Discriminator condition
                 ? new RelationalEntityShaperExpression(EntityType, ValueBufferExpression, true)
                 : this;
+        }
 
         /// <inheritdoc />
         public override EntityShaperExpression Update(Expression valueBufferExpression)

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -298,7 +298,8 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
 
             ((SelectExpression)source1.QueryExpression).ApplyUnion((SelectExpression)source2.QueryExpression, distinct: false);
 
-            return source1;
+            return source1.UpdateShaperExpression(
+                MatchShaperNullabilityForSetOperation(source1.ShaperExpression, source2.ShaperExpression, makeNullable: true));
         }
 
         /// <inheritdoc />
@@ -410,6 +411,8 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
             Check.NotNull(source2, nameof(source2));
 
             ((SelectExpression)source1.QueryExpression).ApplyExcept((SelectExpression)source2.QueryExpression, distinct: true);
+
+            // Since except has result from source1, we don't need to change shaper
             return source1;
         }
 
@@ -588,7 +591,9 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
 
             ((SelectExpression)source1.QueryExpression).ApplyIntersect((SelectExpression)source2.QueryExpression, distinct: true);
 
-            return source1;
+            // For intersect since result comes from both sides, if one of them is non-nullable then both are non-nullable
+            return source1.UpdateShaperExpression(
+                MatchShaperNullabilityForSetOperation(source1.ShaperExpression, source2.ShaperExpression, makeNullable: false));
         }
 
         /// <inheritdoc />
@@ -1195,7 +1200,9 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
             Check.NotNull(source2, nameof(source2));
 
             ((SelectExpression)source1.QueryExpression).ApplyUnion((SelectExpression)source2.QueryExpression, distinct: true);
-            return source1;
+
+            return source1.UpdateShaperExpression(
+                MatchShaperNullabilityForSetOperation(source1.ShaperExpression, source2.ShaperExpression, makeNullable: true));
         }
 
         /// <inheritdoc />
@@ -1602,6 +1609,50 @@ private static void HandleGroupByForAggregate(SelectExpression selectExpression,
             }
         }
 
+        private Expression MatchShaperNullabilityForSetOperation(Expression shaper1, Expression shaper2, bool makeNullable)
+        {
+            switch (shaper1)
+            {
+                case EntityShaperExpression entityShaperExpression1
+                when shaper2 is EntityShaperExpression entityShaperExpression2:
+                    return entityShaperExpression1.IsNullable != entityShaperExpression2.IsNullable
+                        ? entityShaperExpression1.MakeNullable(makeNullable)
+                        : entityShaperExpression1;
+
+                case NewExpression newExpression1
+                when shaper2 is NewExpression newExpression2:
+                    var newArguments = new Expression[newExpression1.Arguments.Count];
+                    for (var i = 0; i < newArguments.Length; i++)
+                    {
+                        newArguments[i] = MatchShaperNullabilityForSetOperation(
+                            newExpression1.Arguments[i], newExpression2.Arguments[i], makeNullable);
+                    }
+
+                    return newExpression1.Update(newArguments);
+
+                case MemberInitExpression memberInitExpression1
+                when shaper2 is MemberInitExpression memberInitExpression2:
+                    var newExpression = (NewExpression)MatchShaperNullabilityForSetOperation(
+                        memberInitExpression1.NewExpression, memberInitExpression2.NewExpression, makeNullable);
+
+                    var memberBindings = new MemberBinding[memberInitExpression1.Bindings.Count];
+                    for (var i = 0; i < memberBindings.Length; i++)
+                    {
+                        var memberAssignment = memberInitExpression1.Bindings[i] as MemberAssignment;
+                        Check.DebugAssert(memberAssignment != null, "Only member assignment bindings are supported");
+
+
+                        memberBindings[i] = memberAssignment.Update(MatchShaperNullabilityForSetOperation(
+                            memberAssignment.Expression, ((MemberAssignment)memberInitExpression2.Bindings[i]).Expression, makeNullable));
+                    }
+
+                    return memberInitExpression1.Update(newExpression, memberBindings);
+
+                default:
+                    return shaper1;
+            }
+        }
+
         private ShapedQueryExpression? AggregateResultShaper(
             ShapedQueryExpression source,
             Expression? projection,

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

@@ -1662,7 +1662,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
                 Check.NotNull(extensionExpression, nameof(extensionExpression));
 
                 return extensionExpression is EntityShaperExpression entityShaper
-                    ? entityShaper.MarkAsNullable()
+                    ? entityShaper.MakeNullable()
                     : base.VisitExtension(extensionExpression);
             }
         }

src/EFCore/Query/EntityShaperExpression.cs

@@ -229,10 +229,19 @@ public virtual EntityShaperExpression WithEntityType([NotNull] IEntityType entit
         ///     Marks this shaper as nullable, indicating that it can shape null entity instances.
         /// </summary>
         /// <returns> This expression if nullability not changed, or an expression with updated nullability. </returns>
+        [Obsolete("Use MakeNullable() instead.")]
         public virtual EntityShaperExpression MarkAsNullable()
-            => !IsNullable
+            => MakeNullable();
+
+        /// <summary>
+        ///     Assigns nullability for this shaper, indicating whether it can shape null entity instances or not.
+        /// </summary>
+        /// <param name="nullable"> A value indicating if the shaper is nullable. </param>
+        /// <returns> This expression if nullability not changed, or an expression with updated nullability. </returns>
+        public virtual EntityShaperExpression MakeNullable(bool nullable = true)
+            => IsNullable != nullable
                 // Marking nullable requires recomputation of materialization condition
-                ? new EntityShaperExpression(EntityType, ValueBufferExpression, true)
+                ? new EntityShaperExpression(EntityType, ValueBufferExpression, nullable)
                 : this;
 
         /// <summary>

src/EFCore/Query/QueryableMethodTranslatingExpressionVisitor.cs

@@ -506,7 +506,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
                 Check.NotNull(extensionExpression, nameof(extensionExpression));
 
                 return extensionExpression is EntityShaperExpression entityShaper
-                    ? entityShaper.MarkAsNullable()
+                    ? entityShaper.MakeNullable()
                     : base.VisitExtension(extensionExpression);
             }
         }