Implement entity equality OrderBy() support

Fixes #15939

src/EFCore/Query/Pipeline/EntityEqualityRewritingExpressionVisitor.cs

@@ -158,6 +158,12 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                     case nameof(Queryable.Contains) when arguments.Count == 2:
                         return VisitContainsMethodCall(methodCallExpression);
 
+                    case nameof(Queryable.OrderBy) when arguments.Count == 2:
+                    case nameof(Queryable.OrderByDescending) when arguments.Count == 2:
+                    case nameof(Queryable.ThenBy) when arguments.Count == 2:
+                    case nameof(Queryable.ThenByDescending) when arguments.Count == 2:
+                        return VisitOrderingMethodCall(methodCallExpression);
+
                     // The following are projecting methods, which flow the entity type from *within* the lambda outside.
                     case nameof(Queryable.Select):
                     case nameof(Queryable.SelectMany):
@@ -207,7 +213,7 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                     || sourceParamType == typeof(IQueryable))   // OfType
                 {
                     // If the method returns the element same type as the source, flow the type information
-                    // (e.g. Where, OrderBy)
+                    // (e.g. Where)
                     if (methodCallExpression.Method.ReturnType.TryGetSequenceType() is Type returnElementType
                         && (returnElementType == sourceElementType || sourceElementType == null))
                     {
@@ -289,6 +295,74 @@ protected virtual Expression VisitContainsMethodCall(MethodCallExpression method
             );
         }
 
+        protected virtual Expression VisitOrderingMethodCall(MethodCallExpression methodCallExpression)
+        {
+            var arguments = methodCallExpression.Arguments;
+            var newSource = Visit(arguments[0]);
+
+            if (!(newSource is EntityReferenceExpression sourceWrapper))
+            {
+                return methodCallExpression.Update(null, new[] { newSource, Unwrap(Visit(arguments[1])) });
+            }
+
+            var newKeySelector = RewriteAndVisitLambda(arguments[1].UnwrapLambdaFromQuote(), sourceWrapper);
+
+            if (!(newKeySelector.Body is EntityReferenceExpression keySelectorWrapper)
+                || !(keySelectorWrapper.EntityType is IEntityType entityType))
+            {
+                return sourceWrapper.Update(
+                    methodCallExpression.Update(null, new[] { Unwrap(newSource), Unwrap(newKeySelector) }));
+            }
+
+            var genericMethodDefinition = methodCallExpression.Method.GetGenericMethodDefinition();
+            var isFirstOrdering =
+                genericMethodDefinition == LinqMethodHelpers.QueryableOrderByMethodInfo
+                || genericMethodDefinition == LinqMethodHelpers.QueryableOrderByDescendingMethodInfo;
+            var isAscending =
+                genericMethodDefinition == LinqMethodHelpers.QueryableOrderByMethodInfo
+                || genericMethodDefinition == LinqMethodHelpers.QueryableThenByMethodInfo;
+
+            var keyProperties = entityType.FindPrimaryKey().Properties;
+            var expression = Unwrap(newSource);
+            var body = Unwrap(newKeySelector.Body);
+            var oldParam = newKeySelector.Parameters.Single();
+
+            foreach (var keyProperty in keyProperties)
+            {
+                var param = Expression.Parameter(entityType.ClrType, "v");
+                var rewrittenKeySelector = Expression.Lambda(
+                    ReplacingExpressionVisitor.Replace(
+                        oldParam, param,
+                        body.CreateEFPropertyExpression(keyProperty, makeNullable: false)),
+                    param);
+
+                var orderingMethodInfo = GetOrderingMethodInfo(isFirstOrdering, isAscending);
+
+                expression = Expression.Call(
+                    orderingMethodInfo.MakeGenericMethod(entityType.ClrType, keyProperty.ClrType),
+                    expression,
+                    rewrittenKeySelector
+                );
+
+                isFirstOrdering = false;
+            }
+
+            return expression;
+
+            static MethodInfo GetOrderingMethodInfo(bool isFirstOrdering, bool isAscending)
+            {
+                if (isFirstOrdering)
+                {
+                    return isAscending
+                        ? LinqMethodHelpers.QueryableOrderByMethodInfo
+                        : LinqMethodHelpers.QueryableOrderByDescendingMethodInfo;
+                }
+                return isAscending
+                    ? LinqMethodHelpers.QueryableThenByMethodInfo
+                    : LinqMethodHelpers.QueryableThenByDescendingMethodInfo;
+            }
+        }
+
         protected virtual Expression VisitSelectMethodCall(MethodCallExpression methodCallExpression)
         {
             var arguments = methodCallExpression.Arguments;

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

@@ -525,6 +525,24 @@ from c in cs.Include(c => c.Orders)
                     where c == null
                     select c.CustomerID);
 
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Entity_equality_orderby(bool isAsync)
+            => AssertQuery<Customer>(
+                isAsync,
+                cs => cs.OrderBy(c => c),
+                entryCount: 91,
+                assertOrder: true);
+
+        [ConditionalTheory]
+        [MemberData(nameof(IsAsyncData))]
+        public virtual Task Entity_equality_orderby_descending_composite_key(bool isAsync)
+            => AssertQuery<OrderDetail>(
+                isAsync,
+                od => od.OrderByDescending(o => o),
+                entryCount: 2155,
+                assertOrder: true);
+
         [ConditionalTheory]
         [MemberData(nameof(IsAsyncData))]
         public virtual Task Null_conditional_simple(bool isAsync)

test/EFCore.Specification.Tests/TestModels/Northwind/Customer.cs

@@ -1,6 +1,7 @@
 // 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;
 using System.Collections.Generic;
 using System.ComponentModel.DataAnnotations.Schema;
 using Microsoft.EntityFrameworkCore.Infrastructure;
@@ -9,7 +10,7 @@
 
 namespace Microsoft.EntityFrameworkCore.TestModels.Northwind
 {
-    public class Customer
+    public class Customer : IComparable<Customer>
     {
         public Customer()
         {
@@ -59,6 +60,9 @@ public override bool Equals(object obj)
 
         public static bool operator !=(Customer left, Customer right) => !Equals(left, right);
 
+        public int CompareTo(Customer other) =>
+            other == null ? 1 : CustomerID.CompareTo(other.CustomerID);
+
         public override int GetHashCode() => CustomerID.GetHashCode();
 
         public override string ToString() => "Customer " + CustomerID;

test/EFCore.Specification.Tests/TestModels/Northwind/OrderDetail.cs

@@ -5,7 +5,7 @@
 
 namespace Microsoft.EntityFrameworkCore.TestModels.Northwind
 {
-    public class OrderDetail
+    public class OrderDetail : IComparable<OrderDetail>
     {
         public int OrderID { get; set; }
         public int ProductID { get; set; }
@@ -34,5 +34,18 @@ public override bool Equals(object obj)
         }
 
         public override int GetHashCode() => HashCode.Combine(OrderID, ProductID);
+
+        public int CompareTo(OrderDetail other)
+        {
+            if (other == null)
+            {
+                return 1;
+            }
+
+            var comp1 = OrderID.CompareTo(other.OrderID);
+            return comp1 == 0
+                ? ProductID.CompareTo(other.ProductID)
+                : comp1;
+        }
     }
 }

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

@@ -283,6 +283,26 @@ FROM [Customers] AS [c]
 WHERE CAST(0 AS bit) = CAST(1 AS bit)");
         }
 
+        public override async Task Entity_equality_orderby(bool isAsync)
+        {
+            await base.Entity_equality_orderby(isAsync);
+
+            AssertSql(
+                @"SELECT [c].[CustomerID], [c].[Address], [c].[City], [c].[CompanyName], [c].[ContactName], [c].[ContactTitle], [c].[Country], [c].[Fax], [c].[Phone], [c].[PostalCode], [c].[Region]
+FROM [Customers] AS [c]
+ORDER BY [c].[CustomerID]");
+        }
+
+        public override async Task Entity_equality_orderby_descending_composite_key(bool isAsync)
+        {
+            await base.Entity_equality_orderby_descending_composite_key(isAsync);
+
+            AssertSql(
+                @"SELECT [o].[OrderID], [o].[ProductID], [o].[Discount], [o].[Quantity], [o].[UnitPrice]
+FROM [Order Details] AS [o]
+ORDER BY [o].[OrderID] DESC, [o].[ProductID] DESC");
+        }
+
         public override async Task Queryable_reprojection(bool isAsync)
         {
             await base.Queryable_reprojection(isAsync);