Change OPENJSON WITH logic

As a temporary representation, always produce OPENJSON with WITH in
translation - including ORDER BY [key]. If that ORDER BY is still
present in post-processing, convert the OPENJSON with WITH to OPENJSON
without WITH.

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -575,7 +575,7 @@ private static ShapedQueryExpression CreateShapedQueryExpression(IEntityType ent
 
     /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateCount(ShapedQueryExpression source, LambdaExpression? predicate)
-        => TranslateAggregateWithPredicate(source, predicate, QueryableMethods.CountWithoutPredicate);
+        => TranslateAggregateWithPredicate(source, predicate, QueryableMethods.CountWithoutPredicate, liftOrderings: false);
 
     /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateDefaultIfEmpty(ShapedQueryExpression source, Expression? defaultValue)
@@ -914,7 +914,7 @@ private SqlExpression CreateJoinPredicate(Expression outerKey, Expression innerK
 
     /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateLongCount(ShapedQueryExpression source, LambdaExpression? predicate)
-        => TranslateAggregateWithPredicate(source, predicate, QueryableMethods.LongCountWithoutPredicate);
+        => TranslateAggregateWithPredicate(source, predicate, QueryableMethods.LongCountWithoutPredicate, liftOrderings: false);
 
     /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateMax(ShapedQueryExpression source, LambdaExpression? selector, Type resultType)
@@ -2377,7 +2377,8 @@ private static Expression MatchShaperNullabilityForSetOperation(Expression shape
     private ShapedQueryExpression? TranslateAggregateWithPredicate(
         ShapedQueryExpression source,
         LambdaExpression? predicate,
-        MethodInfo predicateLessMethodInfo)
+        MethodInfo predicateLessMethodInfo,
+        bool liftOrderings)
     {
         if (predicate != null)
         {
@@ -2396,7 +2397,7 @@ private static Expression MatchShaperNullabilityForSetOperation(Expression shape
             selectExpression.ReplaceProjection(new List<Expression>());
         }
 
-        selectExpression.PrepareForAggregate();
+        selectExpression.PrepareForAggregate(liftOrderings);
         var selector = _sqlExpressionFactory.Fragment("*");
         var methodCall = Expression.Call(
             predicateLessMethodInfo.MakeGenericMethod(selector.Type),

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

@@ -3920,14 +3920,14 @@ public bool IsNonComposedFromSql()
     /// <summary>
     ///     Prepares the <see cref="SelectExpression" /> to apply aggregate operation over it.
     /// </summary>
-    public void PrepareForAggregate()
+    public void PrepareForAggregate(bool liftOrderings = true)
     {
         if (IsDistinct
             || Limit != null
             || Offset != null
             || _groupBy.Count > 0)
         {
-            PushdownIntoSubquery();
+            PushdownIntoSubqueryInternal(liftOrderings);
         }
     }
 

src/EFCore.SqlServer/Query/Internal/SqlServerQueryTranslationPostprocessor.cs

@@ -84,8 +84,9 @@ private sealed class SkipWithoutOrderByInSplitQueryVerifier : ExpressionVisitor
     }
 
     /// <summary>
-    ///     Converts <see cref="SqlServerOpenJsonExpression" /> expressions with WITH (the default) to ordered OPENJSON without WITH when
-    ///     it's detected that the ordering of the original JSON array needs to be preserved (i.e. limit/offset).
+    ///     Converts <see cref="SqlServerOpenJsonExpression" /> expressions with WITH (the default) to OPENJSON without WITH when an
+    ///     ordering still exists on the [key] column, i.e. when the ordering of the original JSON array needs to be preserved
+    ///     (e.g. limit/offset).
     /// </summary>
     private sealed class OpenJsonPostprocessor : ExpressionVisitor
     {
@@ -111,11 +112,20 @@ public Expression Process(Expression expression)
                     return shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
 
                 case SelectExpression
-                    {
-                        Tables: [SqlServerOpenJsonExpression { ColumnInfos: not null } openJsonExpression, ..],
-                        Orderings: []
-                    } selectExpression
-                    when selectExpression.Limit is not null || selectExpression.Offset is not null:
+                {
+                    Tables: [SqlServerOpenJsonExpression { ColumnInfos: not null } openJsonExpression, ..],
+                    Orderings:
+                    [
+                        {
+                            Expression: SqlUnaryExpression
+                            {
+                                OperatorType: ExpressionType.Convert,
+                                Operand: ColumnExpression { Name: "key", Table: var keyColumnTable }
+                            }
+                        }
+                    ]
+                } selectExpression
+                    when keyColumnTable == openJsonExpression:
                 {
                     // Remove the WITH clause from the OPENJSON expression
                     var newOpenJsonExpression = openJsonExpression.Update(
@@ -132,24 +142,10 @@ public Expression Process(Expression expression)
                         selectExpression.Predicate,
                         selectExpression.GroupBy,
                         selectExpression.Having,
-                        Array.Empty<OrderingExpression>(),
+                        selectExpression.Orderings,
                         selectExpression.Limit,
                         selectExpression.Offset);
 
-                    // Add ORDER BY CAST([i].[key] AS int)
-                    newSelectExpression.AppendOrdering(
-                        new OrderingExpression(
-                            _sqlExpressionFactory.Convert(
-                                selectExpression.CreateColumnExpression(
-                                    openJsonExpression,
-                                    "key",
-                                    typeof(string),
-                                    typeMapping: _typeMappingSource.FindMapping("nvarchar(4000)"),
-                                    columnNullable: false),
-                                typeof(int),
-                                _typeMappingSource.FindMapping(typeof(int))),
-                            ascending: true));
-
                     // Record the OPENJSON expression and its projected column(s), along with the store type we just removed from the WITH
                     // clause. Then visit the select expression, adding a cast around the matching ColumnExpressions.
                     // TODO: Need to pass through the type mapping API for converting the JSON value (nvarchar) to the relational store type

src/EFCore.SqlServer/Query/Internal/SqlServerQueryableMethodTranslatingExpressionVisitor.cs

@@ -128,12 +128,6 @@ protected override ShapedQueryExpression TranslateCollection(
         // (i.e. with a columnInfo), which determines the type conversion to apply to the JSON elements coming out.
         // For parameter collections, the element type mapping will only be inferred and applied later (see
         // SqlServerInferredTypeMappingApplier below), at which point the we'll apply it to add the WITH clause.
-
-        // Note also that OPENJSON doesn't guarantee the ordering of the elements coming out. During post-processing, we'll detect whether
-        // order preservation is actually required (e.g. limit/offset are specified), and convert the OPENJSON with WITH to an ordered
-        // OPENJSON without WITH, where we order by the projected 'key' column. Therefore, the OPENJSON with WITH expression may be a
-        // temporary representation until it's replaced in postprocessing.
-
         var openJsonExpression = elementTypeMapping is null
             ? new SqlServerOpenJsonExpression(tableAlias, sqlExpression)
             : new SqlServerOpenJsonExpression(
@@ -148,27 +142,34 @@ protected override ShapedQueryExpression TranslateCollection(
 
         var selectExpression = new SelectExpression(
             openJsonExpression, columnName: "value", columnType: elementClrType, columnTypeMapping: elementTypeMapping, isColumnNullable);
+
+        // OPENJSON doesn't guarantee the ordering of the elements coming out; when using OPENJSON without WITH, a [key] column is returned
+        // with the JSON array's ordering, which we can ORDER BY; this option doesn't exist with OPENJSON with WITH, unfortunately.
+        // However, OPENJSON with WITH has better performance, and also applies JSON-specific conversions we cannot be done otherwise
+        // (e.g. OPENJSON with WITH does base64 decoding for VARBINARY).
+        // Here we generate OPENJSON with WITH, but also add an ordering by [key] - this is a temporary invalid representation.
+        // In SqlServerQueryTranslationPostprocessor, we'll post-process the expression; if the ORDER BY was stripped (e.g. because of
+        // IN, EXISTS or a set operation), we'll just leave the OPENJSON with WITH. If not, we'll convert the OPENJSON with WITH to an
+        // OPENJSON without WITH.
+        // Note that the OPENJSON 'key' column is an nvarchar - we convert it to an int before sorting.
+        selectExpression.AppendOrdering(
+            new OrderingExpression(
+                _sqlExpressionFactory.Convert(
+                    selectExpression.CreateColumnExpression(
+                        openJsonExpression,
+                        "key",
+                        typeof(string),
+                        typeMapping: _typeMappingSource.FindMapping("nvarchar(4000)"),
+                        columnNullable: false),
+                    typeof(int),
+                    _typeMappingSource.FindMapping(typeof(int))),
+                ascending: true));
+
         var shaperExpression = new ProjectionBindingExpression(selectExpression, new ProjectionMember(), elementClrType);
 
         return new ShapedQueryExpression(selectExpression, shaperExpression);
     }
 
-    /// <summary>
-    ///     Consider select expressions whose primary table is a <see cref="SqlServerOpenJsonExpression" /> as naturally ordered, i.e.
-    ///     don't warn if offset/limit is composed on top. The <see cref="SqlServerOpenJsonExpression" /> created in
-    ///     <see cref="TranslateCollection" /> has a WITH clause and indeed doesn't guarantee ordering, but in
-    ///     <see cref="SqlServerQueryTranslationPostprocessor" />we detect cases where ordering is required, remove the WITH clause and
-    ///     inject an ordering by the OPENJSON <c>key</c> column, guaranteeing the natural JSON array ordering.
-    /// </summary>
-    /// <remarks>
-    ///     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.
-    /// </remarks>
-    protected override bool IsOrdered(SelectExpression selectExpression)
-        => base.IsOrdered(selectExpression) || selectExpression.Tables is [SqlServerOpenJsonExpression, ..];
-
     /// <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

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

@@ -811,7 +811,7 @@ public override async Task Parameter_collection_in_subquery_Count_as_compiled_qu
     {
         await base.Parameter_collection_in_subquery_Count_as_compiled_query(async);
 
-        // TODO: the subquery projection contains an extra column which we should remove
+        // TODO: the subquery projection contains extra columns which we should remove
         AssertSql(
 """
 @__ints='[10,111]' (Size = 4000)
@@ -821,7 +821,7 @@ FROM [PrimitiveCollectionsEntity] AS [p]
 WHERE (
     SELECT COUNT(*)
     FROM (
-        SELECT CAST([i].[value] AS int) AS [value], CAST([i].[value] AS int) AS [value0]
+        SELECT CAST([i].[value] AS int) AS [value], CAST([i].[key] AS int) AS [c], CAST([i].[value] AS int) AS [value0]
         FROM OPENJSON(@__ints) AS [i]
         ORDER BY CAST([i].[key] AS int)
         OFFSET 1 ROWS