Stop pushing down to subquery for non-Concat set operations without limit/offset

Closes dotnet#30684

Author: roji

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

@@ -2280,23 +2280,30 @@ private void ApplySetOperation(SetOperationType setOperationType, SelectExpressi
         var entityProjectionValueComparers = new List<ValueComparer>();
         var otherExpressions = new List<SqlExpression>();
 
-        if (select1.Orderings.Count != 0
-            || select1.Limit != null
-            || select1.Offset != null)
+        // Push down into a subquery if limit/offset are defined. If not, any orderings can be discarded as set operations don't preserve
+        // them.
+        // Note that in some databases it may be possible to preserve the internal ordering of the set operands for Concat, but we don't
+        // currently support that.
+        if (select1.Limit != null || select1.Offset != null)
         {
             // If we are pushing down here, we need to make sure to assign unique alias to subquery also.
             var subqueryAlias = GenerateUniqueAlias(_usedAliases, "t");
-            select1.PushdownIntoSubquery();
+            select1.PushdownIntoSubqueryInternal(liftOrderings: false);
             select1._tables[0].Alias = subqueryAlias;
             select1._tableReferences[0].Alias = subqueryAlias;
+        }
+        else
+        {
             select1.ClearOrdering();
         }
 
-        if (select2.Orderings.Count != 0
-            || select2.Limit != null
-            || select2.Offset != null)
+        // Do the same for the other side of the set operation
+        if (select2.Limit != null || select2.Offset != null)
+        {
+            select2.PushdownIntoSubqueryInternal(liftOrderings: false);
+        }
+        else
         {
-            select2.PushdownIntoSubquery();
             select2.ClearOrdering();
         }
 
@@ -3572,7 +3579,11 @@ public Expression AddOuterApply(
     public void PushdownIntoSubquery()
         => PushdownIntoSubqueryInternal();
 
-    private SqlRemappingVisitor PushdownIntoSubqueryInternal()
+    /// <summary>
+    ///     Pushes down the <see cref="SelectExpression" /> into a subquery.
+    /// </summary>
+    /// <param name="liftOrderings">Whether orderings on the query should be lifted out of the subquery.</param>
+    private SqlRemappingVisitor PushdownIntoSubqueryInternal(bool liftOrderings = true)
     {
         var subqueryAlias = GenerateUniqueAlias(_usedAliases, "t");
         var subquery = new SelectExpression(
@@ -3736,13 +3747,14 @@ private SqlRemappingVisitor PushdownIntoSubqueryInternal()
         foreach (var ordering in subquery._orderings)
         {
             var orderingExpression = ordering.Expression;
-            if (projectionMap.TryGetValue(orderingExpression, out var outerColumn))
+            if (liftOrderings && projectionMap.TryGetValue(orderingExpression, out var outerColumn))
             {
                 _orderings.Add(ordering.Update(outerColumn));
             }
-            else if (!IsDistinct
-                     && GroupBy.Count == 0
-                     || GroupBy.Contains(orderingExpression))
+            else if (liftOrderings
+                     && (!IsDistinct
+                         && GroupBy.Count == 0
+                         || GroupBy.Contains(orderingExpression)))
             {
                 _orderings.Add(
                     ordering.Update(

test/EFCore.Specification.Tests/Query/NorthwindSetOperationsQueryTestBase.cs

@@ -761,6 +761,38 @@ public virtual Task Union_over_scalarsubquery_scalarsubquery(bool async)
             ss => ss.Set<Order>().Select(o => o.OrderDetails.Count())
                 .Union(ss.Set<Order>().Select(o => o.OrderDetails.Count())));
 
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Union_over_OrderBy_Take1(bool async)
+        => AssertQueryScalar(
+            async,
+            ss => ss.Set<Order>().OrderBy(o => o.OrderDate).Take(5).Select(o => o.OrderID)
+                .Union(ss.Set<Order>().Select(o => o.OrderID)));
+
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Union_over_OrderBy_without_Skip_Take1(bool async)
+        => AssertQueryScalar(
+            async,
+            ss => ss.Set<Order>().OrderBy(o => o.OrderDate).Select(o => o.OrderID)
+                .Union(ss.Set<Order>().Select(o => o.OrderID)));
+
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Union_over_OrderBy_Take2(bool async)
+        => AssertQueryScalar(
+            async,
+            ss => ss.Set<Order>().Select(o => o.OrderID)
+                .Union(ss.Set<Order>().OrderBy(o => o.OrderDate).Take(5).Select(o => o.OrderID)));
+
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Union_over_OrderBy_without_Skip_Take2(bool async)
+        => AssertQueryScalar(
+            async,
+            ss => ss.Set<Order>().Select(o => o.OrderID)
+                .Union(ss.Set<Order>().OrderBy(o => o.OrderDate).Select(o => o.OrderID)));
+
     [ConditionalTheory]
     [MemberData(nameof(IsAsyncData))]
     public virtual Task OrderBy_Take_Union(bool async)

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

@@ -1020,6 +1020,74 @@ FROM [Orders] AS [o1]
 """);
     }
 
+    public override async Task Union_over_OrderBy_Take1(bool async)
+    {
+        await base.Union_over_OrderBy_Take1(async);
+
+        AssertSql(
+"""
+@__p_0='5'
+
+SELECT [t].[OrderID]
+FROM (
+    SELECT TOP(@__p_0) [o].[OrderID]
+    FROM [Orders] AS [o]
+    ORDER BY [o].[OrderDate]
+) AS [t]
+UNION
+SELECT [o0].[OrderID]
+FROM [Orders] AS [o0]
+""");
+    }
+
+    public override async Task Union_over_OrderBy_without_Skip_Take1(bool async)
+    {
+        await base.Union_over_OrderBy_without_Skip_Take1(async);
+
+        AssertSql(
+"""
+SELECT [o].[OrderID]
+FROM [Orders] AS [o]
+UNION
+SELECT [o0].[OrderID]
+FROM [Orders] AS [o0]
+""");
+    }
+
+    public override async Task Union_over_OrderBy_Take2(bool async)
+    {
+        await base.Union_over_OrderBy_Take2(async);
+
+        AssertSql(
+"""
+@__p_0='5'
+
+SELECT [o].[OrderID]
+FROM [Orders] AS [o]
+UNION
+SELECT [t0].[OrderID]
+FROM (
+    SELECT TOP(@__p_0) [o0].[OrderID]
+    FROM [Orders] AS [o0]
+    ORDER BY [o0].[OrderDate]
+) AS [t0]
+""");
+    }
+
+    public override async Task Union_over_OrderBy_without_Skip_Take2(bool async)
+    {
+        await base.Union_over_OrderBy_without_Skip_Take2(async);
+
+        AssertSql(
+"""
+SELECT [o].[OrderID]
+FROM [Orders] AS [o]
+UNION
+SELECT [o0].[OrderID]
+FROM [Orders] AS [o0]
+""");
+    }
+
     public override async Task OrderBy_Take_Union(bool async)
     {
         await base.OrderBy_Take_Union(async);