Generate OPENJSON with WITH unless ordering is required

Part of dotnet#13617

src/EFCore.Relational/Query/QuerySqlGenerator.cs

@@ -697,8 +697,7 @@ string GetUniqueParameterName(string currentName)
     /// <inheritdoc />
     protected override Expression VisitOrdering(OrderingExpression orderingExpression)
     {
-        if (orderingExpression.Expression is SqlConstantExpression
-            || orderingExpression.Expression is SqlParameterExpression)
+        if (orderingExpression.Expression is SqlConstantExpression or SqlParameterExpression)
         {
             _relationalCommandBuilder.Append("(SELECT 1)");
         }

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.Helper.cs

@@ -1014,6 +1014,9 @@ private sealed class CloningExpressionVisitor : ExpressionVisitor
                     return newTpcTable;
                 }
 
+                case IClonableTableExpressionBase cloneable:
+                    return cloneable.Clone();
+
                 case TableValuedFunctionExpression tableValuedFunctionExpression:
                 {
                     var newArguments = new SqlExpression[tableValuedFunctionExpression.Arguments.Count];
@@ -1036,9 +1039,6 @@ private sealed class CloningExpressionVisitor : ExpressionVisitor
                     return newTableValuedFunctionExpression;
                 }
 
-                case IClonableTableExpressionBase cloneable:
-                    return cloneable.Clone();
-
                 // join and set operations are fine, because they contain other TableExpressionBases inside, that will get cloned
                 // and therefore set expression's Update function will generate a new instance.
                 case JoinExpressionBase or SetOperationBase:

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

@@ -3919,14 +3919,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);
         }
     }
 
@@ -4664,10 +4664,6 @@ protected override void Print(ExpressionPrinter expressionPrinter)
             expressionPrinter.AppendLine().Append("ORDER BY ");
             expressionPrinter.VisitCollection(Orderings);
         }
-        else if (Offset != null)
-        {
-            expressionPrinter.AppendLine().Append("ORDER BY (SELECT 1)");
-        }
 
         if (Offset != null)
         {

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

@@ -116,6 +116,7 @@ public virtual ValuesExpression Update(IReadOnlyList<RowValueExpression> rowValu
     protected override TableExpressionBase CreateWithAnnotations(IEnumerable<IAnnotation> annotations)
         => new ValuesExpression(Alias, RowValues, ColumnNames, annotations);
 
+    // TODO: Deep clone, see #30982
     /// <inheritdoc />
     public virtual TableExpressionBase Clone()
         => CreateWithAnnotations(GetAnnotations());

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

@@ -20,7 +20,7 @@ namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal;
 ///         doing so can result in application failures when updating to a new Entity Framework Core release.
 ///     </para>
 /// </remarks>
-public class SqlServerOpenJsonExpression : TableValuedFunctionExpression
+public class SqlServerOpenJsonExpression : TableValuedFunctionExpression, IClonableTableExpressionBase
 {
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -80,6 +80,26 @@ public virtual SqlExpression JsonExpression
             ? this
             : new SqlServerOpenJsonExpression(Alias, jsonExpression, path, columnInfos);
 
+
+    /// <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>
+    // TODO: Deep clone, see #30982
+    public virtual TableExpressionBase Clone()
+    {
+        var clone = new SqlServerOpenJsonExpression(Alias, JsonExpression, Path, ColumnInfos);
+
+        foreach (var annotation in GetAnnotations())
+        {
+            clone.AddAnnotation(annotation.Name, annotation.Value);
+        }
+
+        return clone;
+    }
+
     /// <inheritdoc />
     protected override void Print(ExpressionPrinter expressionPrinter)
     {
@@ -145,5 +165,5 @@ public override int GetHashCode()
     ///     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 readonly record struct ColumnInfo(string Name, string? StoreType, string? Path = null, bool AsJson = false);
+    public readonly record struct ColumnInfo(string Name, string StoreType, string? Path = null, bool AsJson = false);
 }

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

@@ -15,8 +15,8 @@ namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal;
 /// </summary>
 public class SqlServerQueryTranslationPostprocessor : RelationalQueryTranslationPostprocessor
 {
-    private readonly SkipWithoutOrderByInSplitQueryVerifyingExpressionVisitor
-        _skipWithoutOrderByInSplitQueryVerifyingExpressionVisitor = new();
+    private readonly OpenJsonPostprocessor _openJsonPostprocessor;
+    private readonly SkipWithoutOrderByInSplitQueryVerifier _skipWithoutOrderByInSplitQueryVerifier = new();
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -27,9 +27,11 @@ private readonly SkipWithoutOrderByInSplitQueryVerifyingExpressionVisitor
     public SqlServerQueryTranslationPostprocessor(
         QueryTranslationPostprocessorDependencies dependencies,
         RelationalQueryTranslationPostprocessorDependencies relationalDependencies,
-        QueryCompilationContext queryCompilationContext)
+        QueryCompilationContext queryCompilationContext,
+        IRelationalTypeMappingSource typeMappingSource)
         : base(dependencies, relationalDependencies, queryCompilationContext)
     {
+        _openJsonPostprocessor = new(typeMappingSource, relationalDependencies.SqlExpressionFactory);
     }
 
     /// <summary>
@@ -40,14 +42,15 @@ private readonly SkipWithoutOrderByInSplitQueryVerifyingExpressionVisitor
     /// </summary>
     public override Expression Process(Expression query)
     {
-        var result = base.Process(query);
+        query = base.Process(query);
 
-        _skipWithoutOrderByInSplitQueryVerifyingExpressionVisitor.Visit(result);
+        query = _openJsonPostprocessor.Process(query);
+        _skipWithoutOrderByInSplitQueryVerifier.Visit(query);
 
-        return result;
+        return query;
     }
 
-    private sealed class SkipWithoutOrderByInSplitQueryVerifyingExpressionVisitor : ExpressionVisitor
+    private sealed class SkipWithoutOrderByInSplitQueryVerifier : ExpressionVisitor
     {
         [return: NotNullIfNotNull("expression")]
         public override Expression? Visit(Expression? expression)
@@ -68,9 +71,7 @@ private sealed class SkipWithoutOrderByInSplitQueryVerifyingExpressionVisitor :
 
                     return relationalSplitCollectionShaperExpression;
 
-                case SelectExpression selectExpression
-                    when selectExpression.Offset != null
-                    && selectExpression.Orderings.Count == 0:
+                case SelectExpression { Offset: not null, Orderings.Count: 0 }:
                     throw new InvalidOperationException(SqlServerStrings.SplitQueryOffsetWithoutOrderBy);
 
                 case NonQueryExpression nonQueryExpression:
@@ -81,4 +82,103 @@ private sealed class SkipWithoutOrderByInSplitQueryVerifyingExpressionVisitor :
             }
         }
     }
+
+    /// <summary>
+    ///     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
+    {
+        private readonly IRelationalTypeMappingSource _typeMappingSource;
+        private readonly ISqlExpressionFactory _sqlExpressionFactory;
+        private readonly Dictionary<(SqlServerOpenJsonExpression, string), RelationalTypeMapping> _castsToApply = new();
+
+        public OpenJsonPostprocessor(IRelationalTypeMappingSource typeMappingSource, ISqlExpressionFactory sqlExpressionFactory)
+            => (_typeMappingSource, _sqlExpressionFactory) = (typeMappingSource, sqlExpressionFactory);
+
+        public Expression Process(Expression expression)
+        {
+            _castsToApply.Clear();
+            return Visit(expression);
+        }
+
+        [return: NotNullIfNotNull("expression")]
+        public override Expression? Visit(Expression? expression)
+        {
+            switch (expression)
+            {
+                case ShapedQueryExpression shapedQueryExpression:
+                    return shapedQueryExpression.UpdateQueryExpression(Visit(shapedQueryExpression.QueryExpression));
+
+                case SelectExpression
+                {
+                    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(
+                        openJsonExpression.JsonExpression,
+                        openJsonExpression.Path,
+                        columnInfos: null);
+
+                    var newTables = selectExpression.Tables.ToArray();
+                    newTables[0] = newOpenJsonExpression;
+
+                    var newSelectExpression = selectExpression.Update(
+                        selectExpression.Projection,
+                        newTables,
+                        selectExpression.Predicate,
+                        selectExpression.GroupBy,
+                        selectExpression.Having,
+                        selectExpression.Orderings,
+                        selectExpression.Limit,
+                        selectExpression.Offset);
+
+                    // 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
+                    // (e.g. datetime2), see #30677
+                    foreach (var column in openJsonExpression.ColumnInfos)
+                    {
+                        var typeMapping = _typeMappingSource.FindMapping(column.StoreType);
+                        Check.DebugAssert(
+                            typeMapping is not null,
+                            $"Could not find mapping for store type {column.StoreType} when converting OPENJSON/WITH");
+
+                        _castsToApply.Add((newOpenJsonExpression, column.Name), typeMapping);
+                    }
+
+                    var result = base.Visit(newSelectExpression);
+
+                    foreach (var column in openJsonExpression.ColumnInfos)
+                    {
+                        _castsToApply.Remove((newOpenJsonExpression, column.Name));
+                    }
+
+                    return result;
+                }
+
+                case ColumnExpression { Table: SqlServerOpenJsonExpression openJsonTable, Name: var name } columnExpression
+                    when _castsToApply.TryGetValue((openJsonTable, name), out var typeMapping):
+                {
+                    return _sqlExpressionFactory.Convert(columnExpression, columnExpression.Type, typeMapping);
+                }
+
+                default:
+                    return base.Visit(expression);
+            }
+        }
+    }
 }

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

@@ -11,6 +11,8 @@ namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal;
 /// </summary>
 public class SqlServerQueryTranslationPostprocessorFactory : IQueryTranslationPostprocessorFactory
 {
+    private readonly IRelationalTypeMappingSource _typeMappingSource;
+
     /// <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
@@ -19,10 +21,12 @@ public class SqlServerQueryTranslationPostprocessorFactory : IQueryTranslationPo
     /// </summary>
     public SqlServerQueryTranslationPostprocessorFactory(
         QueryTranslationPostprocessorDependencies dependencies,
-        RelationalQueryTranslationPostprocessorDependencies relationalDependencies)
+        RelationalQueryTranslationPostprocessorDependencies relationalDependencies,
+        IRelationalTypeMappingSource typeMappingSource)
     {
         Dependencies = dependencies;
         RelationalDependencies = relationalDependencies;
+        _typeMappingSource = typeMappingSource;
     }
 
     /// <summary>
@@ -42,8 +46,5 @@ public class SqlServerQueryTranslationPostprocessorFactory : IQueryTranslationPo
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     public virtual QueryTranslationPostprocessor Create(QueryCompilationContext queryCompilationContext)
-        => new SqlServerQueryTranslationPostprocessor(
-            Dependencies,
-            RelationalDependencies,
-            queryCompilationContext);
+        => new SqlServerQueryTranslationPostprocessor(Dependencies, RelationalDependencies, queryCompilationContext, _typeMappingSource);
 }