[release/11.0] Translate set operations over null-checked subqueries

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

@@ -16,6 +16,8 @@ public class RelationalProjectionBindingExpressionVisitor : ExpressionVisitor
 {
     private static readonly MethodInfo GetParameterValueMethodInfo
         = typeof(RelationalProjectionBindingExpressionVisitor).GetTypeInfo().GetDeclaredMethod(nameof(GetParameterValue))!;
+    private static readonly bool UseOldBehavior38838
+        = AppContext.TryGetSwitch("Microsoft.EntityFrameworkCore.Issue38838", out var enabled) && enabled;
 
     private readonly RelationalQueryableMethodTranslatingExpressionVisitor _queryableMethodTranslatingExpressionVisitor;
     private readonly RelationalSqlTranslatingExpressionVisitor _sqlTranslator;
@@ -48,6 +50,62 @@ public RelationalProjectionBindingExpressionVisitor(
         _selectExpression = null!;
     }
 
+    private sealed class MarkerNullCheckSimplifyingExpressionVisitor : ExpressionVisitor
+    {
+        protected override Expression VisitBinary(BinaryExpression node)
+        {
+            if (node is { NodeType: ExpressionType.Equal or ExpressionType.NotEqual, Method: null }
+                && ((IsNull(node.Left) && TryGetNullCheck(node.Right, out var nullCheck))
+                    || (IsNull(node.Right) && TryGetNullCheck(node.Left, out nullCheck))))
+            {
+                nullCheck = Visit(nullCheck);
+                return node.NodeType == ExpressionType.Equal
+                    ? nullCheck
+                    : Expression.Not(nullCheck);
+            }
+
+            return base.VisitBinary(node);
+        }
+
+        private static bool IsNull(Expression expression)
+            => expression is ConstantExpression { Value: null }
+                or DefaultExpression { Type.IsValueType: false };
+
+        private static bool TryGetNullCheck(Expression expression, [NotNullWhen(true)] out Expression? nullCheck)
+        {
+            expression = expression.UnwrapTypeConversion(out _);
+            if (expression is ConditionalExpression
+                {
+                    Test: var test,
+                    IfTrue: var ifTrue,
+                    IfFalse: NewExpression or MemberInitExpression
+                }
+                && IsNull(ifTrue)
+                && IsMarkerNullCheck(test))
+            {
+                nullCheck = test;
+                return true;
+            }
+
+            nullCheck = null;
+            return false;
+        }
+
+        private static bool IsMarkerNullCheck(Expression expression)
+        {
+            expression = expression.UnwrapTypeConversion(out _);
+            return expression is BinaryExpression
+                {
+                    NodeType: ExpressionType.Equal,
+                    Method: null,
+                    Left: var left,
+                    Right: var right
+                }
+                && ((IsNull(left) && right.UnwrapTypeConversion(out _) is ProjectionBindingExpression)
+                    || (IsNull(right) && left.UnwrapTypeConversion(out _) is ProjectionBindingExpression));
+        }
+    }
+
     /// <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
@@ -82,6 +140,7 @@ public virtual Expression Translate(SelectExpression selectExpression, Expressio
 
             _selectExpression.ReplaceProjection(_clientProjections);
             _clientProjections.Clear();
+            _projectionMapping.Clear();
         }
         else
         {
@@ -97,6 +156,36 @@ public virtual Expression Translate(SelectExpression selectExpression, Expressio
         return result;
     }
 
+    internal virtual Expression? TryTranslateToServerProjection(SelectExpression selectExpression, Expression expression)
+    {
+        if (UseOldBehavior38838)
+        {
+            return null;
+        }
+
+        _selectExpression = selectExpression;
+        _indexBasedBinding = false;
+        _rootIsTransparentIdentifier = IsTransparentIdentifierProjection(expression);
+        _projectionMembers.Push(new ProjectionMember());
+
+        expression = new MarkerNullCheckSimplifyingExpressionVisitor().Visit(expression);
+        var result = Visit(expression);
+        if (result == QueryCompilationContext.NotTranslatedExpression)
+        {
+            result = null;
+        }
+        else
+        {
+            _selectExpression.ReplaceProjection(_projectionMapping);
+            result = MatchTypes(result, expression.Type);
+        }
+        _selectExpression = null!;
+        _projectionMapping.Clear();
+        _projectionMembers.Clear();
+
+        return result;
+    }
+
     /// <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.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -585,6 +585,9 @@ protected override ShapedQueryExpression TranslateCast(ShapedQueryExpression sou
     /// <inheritdoc />
     protected override ShapedQueryExpression TranslateConcat(ShapedQueryExpression source1, ShapedQueryExpression source2)
     {
+        source1 = TranslateSetOperationOperand(source1);
+        source2 = TranslateSetOperationOperand(source2);
+
         ((SelectExpression)source1.QueryExpression).ApplyUnion((SelectExpression)source2.QueryExpression, distinct: false);
 
         return source1.UpdateShaperExpression(
@@ -719,6 +722,9 @@ protected override ShapedQueryExpression TranslateDistinct(ShapedQueryExpression
     /// <inheritdoc />
     protected override ShapedQueryExpression TranslateExcept(ShapedQueryExpression source1, ShapedQueryExpression source2)
     {
+        source1 = TranslateSetOperationOperand(source1);
+        source2 = TranslateSetOperationOperand(source2);
+
         ((SelectExpression)source1.QueryExpression).ApplyExcept((SelectExpression)source2.QueryExpression, distinct: true);
 
         // Since except has result from source1, we don't need to change shaper
@@ -884,6 +890,9 @@ protected override ShapedQueryExpression TranslateExcept(ShapedQueryExpression s
     /// <inheritdoc />
     protected override ShapedQueryExpression TranslateIntersect(ShapedQueryExpression source1, ShapedQueryExpression source2)
     {
+        source1 = TranslateSetOperationOperand(source1);
+        source2 = TranslateSetOperationOperand(source2);
+
         ((SelectExpression)source1.QueryExpression).ApplyIntersect((SelectExpression)source2.QueryExpression, distinct: true);
 
         // For intersect since result comes from both sides, if one of them is non-nullable then both are non-nullable
@@ -1592,12 +1601,25 @@ private void ApplyLimit(SelectExpression selectExpression, SqlExpression limit)
     /// <inheritdoc />
     protected override ShapedQueryExpression TranslateUnion(ShapedQueryExpression source1, ShapedQueryExpression source2)
     {
+        source1 = TranslateSetOperationOperand(source1);
+        source2 = TranslateSetOperationOperand(source2);
+
         ((SelectExpression)source1.QueryExpression).ApplyUnion((SelectExpression)source2.QueryExpression, distinct: true);
 
         return source1.UpdateShaperExpression(
             MatchShaperNullabilityForSetOperation(source1.ShaperExpression, source2.ShaperExpression, makeNullable: true));
     }
 
+    private ShapedQueryExpression TranslateSetOperationOperand(ShapedQueryExpression source)
+    {
+        var selectExpression = (SelectExpression)source.QueryExpression;
+        return selectExpression.HasClientProjections
+            && _projectionBindingExpressionVisitor.TryTranslateToServerProjection(selectExpression, source.ShaperExpression)
+            is { } serverShaper
+                ? source.UpdateShaperExpression(serverShaper)
+                : source;
+    }
+
     /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateWhere(ShapedQueryExpression source, LambdaExpression predicate)
     {

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

@@ -39,6 +39,8 @@ public sealed partial class SelectExpression : TableExpressionBase
     private readonly SqlAliasManager _sqlAliasManager;
 
     internal bool IsMutable { get; private set; } = true;
+    internal bool HasClientProjections
+        => _clientProjections.Count > 0;
     private Dictionary<ProjectionMember, Expression> _projectionMapping = [];
     private List<Expression> _clientProjections = [];
     private readonly List<string?> _aliasForClientProjections = [];

test/EFCore.Relational.Specification.Tests/Query/AdHocMiscellaneousQueryRelationalTestBase.cs

@@ -1004,25 +1004,32 @@ public virtual async Task Union_of_two_leftjoin_nonentity()
             using var context = contextFactory.CreateDbContext();
 
             var categories = context.Requests
-                .GroupBy(r => r.PickupStatusId, (k, els) => new { pickupStatusId = k, Count = els.Count() });
+                .GroupBy(
+                    r => r.PickupStatusId,
+                    (k, els) => new Context30915.CountDto30915
+                    {
+                        PickupStatusId = k,
+                        Count = els.Count()
+                    });
 
             var first = from s in context.Statuses
-                        join c in categories on s.PickupStatusId equals c.pickupStatusId into g
+                        join c in categories on s.PickupStatusId equals c.PickupStatusId into g
                         from countInfo in g.DefaultIfEmpty()
                         select new { s.PickupStatusId, Count = countInfo == null ? 0 : countInfo.Count };
 
             var second = from s in context.Statuses
-                         join c in categories on s.PickupStatusId equals c.pickupStatusId into g
+                         join c in categories on s.PickupStatusId equals c.PickupStatusId into g
                          from countInfo in g.DefaultIfEmpty()
                          select new { s.PickupStatusId, Count = countInfo == null ? 0 : countInfo.Count };
 
             var query = first.Union(second);
 
-            // The client-side null-check projection forces a client projection on each operand,
-            // which then can't participate in the set operation.
-            var ex = await Assert.ThrowsAsync<InvalidOperationException>(() => query.ToListAsync());
-            Assert.Contains("Unable to translate set operation", ex.Message);
-            // #30915 TODO: currently throws on base; flip to assert results if/when fixed.
+            var result = await query.OrderBy(e => e.PickupStatusId).ToListAsync();
+
+            Assert.Equal(3, result.Count);
+            Assert.Equal((1, 2), (result[0].PickupStatusId, result[0].Count));
+            Assert.Equal((2, 0), (result[1].PickupStatusId, result[1].Count));
+            Assert.Equal((3, 1), (result[2].PickupStatusId, result[2].Count));
         }
 
         [Fact]

test/EFCore.Relational.Specification.Tests/Query/NorthwindSetOperationsQueryRelationalTestBase.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.EntityFrameworkCore.TestModels.Northwind;
+
 namespace Microsoft.EntityFrameworkCore.Query;
 
 public abstract class NorthwindSetOperationsQueryRelationalTestBase<TFixture>(TFixture fixture)
@@ -22,4 +24,74 @@ public override async Task Collection_projection_before_set_operation_fails(bool
 
         Assert.Equal(RelationalStrings.SetOperationsNotAllowedAfterClientEvaluation, message);
     }
+
+    [Theory, MemberData(nameof(IsAsyncData))]
+    public virtual async Task Set_operations_over_null_checked_to_one_non_entity_subquery(bool async)
+    {
+        await AssertQuery(
+            async,
+            ss =>
+            {
+                var withOrders =
+                    from c in ss.Set<Customer>()
+                    let latest = c.Orders
+                        .OrderByDescending(o => o.OrderDate)
+                        .ThenByDescending(o => o.OrderID)
+                        .Select(o => new { o.OrderID, o.OrderDate })
+                        .FirstOrDefault()
+                    select new
+                    {
+                        c.CustomerID,
+                        c.City,
+                        LatestOrderID = latest != null ? latest.OrderID : 0,
+                        LatestOrderDate = latest != null ? latest.OrderDate : null
+                    };
+
+                var neverOrdered = ss.Set<Customer>()
+                    .Where(c => !c.Orders.Any())
+                    .Select(c => new
+                    {
+                        c.CustomerID,
+                        c.City,
+                        LatestOrderID = 0,
+                        LatestOrderDate = (DateTime?)null
+                    });
+
+                return withOrders.Concat(neverOrdered);
+            },
+            elementSorter: e => (e.CustomerID, e.LatestOrderID));
+
+        await AssertQuery(
+            async,
+            ss =>
+            {
+                var withOrders =
+                    from c in ss.Set<Customer>()
+                    let latest = c.Orders
+                        .OrderByDescending(o => o.OrderDate)
+                        .ThenByDescending(o => o.OrderID)
+                        .Select(o => new { o.OrderID, o.OrderDate })
+                        .FirstOrDefault()
+                    select new
+                    {
+                        c.CustomerID,
+                        c.City,
+                        LatestOrderID = latest == null ? 0 : latest.OrderID,
+                        LatestOrderDate = latest == null ? null : latest.OrderDate
+                    };
+
+                var neverOrdered = ss.Set<Customer>()
+                    .Where(c => !c.Orders.Any())
+                    .Select(c => new
+                    {
+                        c.CustomerID,
+                        c.City,
+                        LatestOrderID = 0,
+                        LatestOrderDate = (DateTime?)null
+                    });
+
+                return neverOrdered.Union(withOrders);
+            },
+            elementSorter: e => (e.CustomerID, e.LatestOrderID));
+    }
 }

test/EFCore.SqlServer.FunctionalTests/Query/AdHocMiscellaneousQuerySqlServerTest.cs

@@ -3237,7 +3237,34 @@ public override async Task Union_of_two_leftjoin_nonentity()
     {
         await base.Union_of_two_leftjoin_nonentity();
 
-        AssertSql();
+        AssertSql(
+            """
+SELECT [u].[PickupStatusId], [u].[Count]
+FROM (
+    SELECT [s].[PickupStatusId], CASE
+        WHEN [r0].[marker] IS NULL THEN 0
+        ELSE [r0].[Count]
+    END AS [Count]
+    FROM [Statuses] AS [s]
+    LEFT JOIN (
+        SELECT [r].[PickupStatusId], COUNT(*) AS [Count], 1 AS [marker]
+        FROM [Requests] AS [r]
+        GROUP BY [r].[PickupStatusId]
+    ) AS [r0] ON [s].[PickupStatusId] = [r0].[PickupStatusId]
+    UNION
+    SELECT [s0].[PickupStatusId], CASE
+        WHEN [r2].[marker] IS NULL THEN 0
+        ELSE [r2].[Count]
+    END AS [Count]
+    FROM [Statuses] AS [s0]
+    LEFT JOIN (
+        SELECT [r1].[PickupStatusId], COUNT(*) AS [Count], 1 AS [marker]
+        FROM [Requests] AS [r1]
+        GROUP BY [r1].[PickupStatusId]
+    ) AS [r2] ON [s0].[PickupStatusId] = [r2].[PickupStatusId]
+) AS [u]
+ORDER BY [u].[PickupStatusId]
+""");
     }
 
     public override async Task OrderBy_member_of_nullable_projection()