Fix to #28648 - Json/Query: translate element access of a json array

Converting indexer over list/array into ElementAt, so that nav expansion understands it and can perform pushown and inject MaterializeCollectionNavigation expression where necessary.
In translation phase we recognize the pattern that nav expansion creates and if the root is JsonQueryExpression, we apply collection index over it.
JsonQueryExpression path segment now consists of two components - string representing JSON property name and SqlExpression representing collection index (it can be constant, parameter or any arbitrary expression that resolves to int)

Deduplication is heavily restricted currently - we only de-duplicate projections whose additional path consists of JSON property accesses only (no collection indexes allowed).
All queries projecting entities that need JSON array access must be set to NoTracking (for now). This is because we don't flow those collection index values into shaper. Instead, the ordinal keys are filled with dummy values, which prohibits us from doing proper change tracking.

Fixes #28648

src/EFCore.Relational/Properties/RelationalStrings.resx

@@ -950,6 +950,9 @@
   <data name="PendingAmbientTransaction" xml:space="preserve">
     <value>This connection was used with an ambient transaction. The original ambient transaction needs to be completed before this connection can be used outside of it.</value>
   </data>
+  <data name="ProjectingJsonCollectionElementRequiresNoTracking" xml:space="preserve">
+    <value>The query projects an entity mapped to JSON and accesses a JSON collection element. Such queries require 'AsNoTracking' option, even when the parent entity is projected.</value>
+  </data>
   <data name="ProjectionMappingCountMismatch" xml:space="preserve">
     <value>Unable to translate set operations when both sides don't assign values to the same properties in the nominal type. Please make sure that the same properties are included on both sides, and consider assigning default values if a property doesn't require a specific value.</value>
   </data>

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

@@ -0,0 +1,57 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+
+namespace Microsoft.EntityFrameworkCore.Query.Internal;
+
+/// <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>
+internal class CollectionIndexerToElementAtNormalizingExpressionVisitor : ExpressionVisitor
+{
+    protected override Expression VisitMethodCall(MethodCallExpression methodCallExpression)
+    {
+        // Convert list[x] to list.ElementAt(x)
+        if (methodCallExpression.Method is { Name: "get_Item", IsStatic: false, DeclaringType: { IsGenericType: true } declaringType }
+            && declaringType.GetGenericTypeDefinition() == typeof(List<>))
+        {
+            var source = Visit(methodCallExpression.Object!);
+            var index = Visit(methodCallExpression.Arguments[0]);
+            var sourceTypeArgument = source.Type.GetSequenceType();
+
+            return Expression.Call(
+                QueryableMethods.ElementAt.MakeGenericMethod(sourceTypeArgument),
+                    Expression.Call(
+                        QueryableMethods.AsQueryable.MakeGenericMethod(sourceTypeArgument),
+                        source),
+                    index);
+        }
+
+        return base.VisitMethodCall(methodCallExpression);
+    }
+
+    protected override Expression VisitBinary(BinaryExpression binaryExpression)
+    {
+        // Convert array[x] to list.ElementAt(x)
+        if (binaryExpression.NodeType == ExpressionType.ArrayIndex
+            && binaryExpression.Left.Type != typeof(byte[]))
+        {
+            var source = Visit(binaryExpression.Left);
+            var index = Visit(binaryExpression.Right);
+            var sourceTypeArgument = source.Type.GetSequenceType();
+
+            return Expression.Call(
+                QueryableMethods.ElementAt.MakeGenericMethod(sourceTypeArgument),
+                    Expression.Call(
+                        QueryableMethods.AsQueryable.MakeGenericMethod(sourceTypeArgument),
+                        source),
+                    index);
+        }
+
+        return base.VisitBinary(binaryExpression);
+    }
+}

src/EFCore.Relational/Query/JsonQueryExpression.cs

@@ -161,6 +161,32 @@ public virtual JsonQueryExpression BindNavigation(INavigation navigation)
             IsNullable || !navigation.ForeignKey.IsRequiredDependent);
     }
 
+    /// <summary>
+    ///     Binds a collection element access with this JSON query expression to get the SQL representation.
+    /// </summary>
+    /// <param name="collectionIndexExpression">The collection index to bind.</param>
+    public virtual JsonQueryExpression BindCollectionElement(SqlExpression collectionIndexExpression)
+    {
+        // this needs to be changed IF JsonQueryExpression will also be used for collection of primitives
+        // see issue #28688
+        Debug.Assert(
+            Path.Last().ArrayIndex == null,
+            "Already accessing JSON array element.");
+
+        var newPath = Path.ToList();
+        newPath.Add(new PathSegment(collectionIndexExpression));
+
+        return new JsonQueryExpression(
+            EntityType,
+            JsonColumn,
+            _keyPropertyMap,
+            newPath,
+            EntityType.ClrType,
+            collection: false,
+            // TODO: is this the right nullable?
+            nullable: true);
+    }
+
     /// <summary>
     ///     Makes this JSON query expression nullable.
     /// </summary>

src/EFCore.Relational/Query/PathSegment.cs

@@ -7,44 +7,61 @@ namespace Microsoft.EntityFrameworkCore.Query;
 
 /// <summary>
 ///     <para>
-///         A class representing a component of JSON path used in <see cref="JsonQueryExpression" /> or <see cref="JsonScalarExpression" />.
+///         A struct representing a component of JSON path used in <see cref="JsonQueryExpression" /> or <see cref="JsonScalarExpression" />.
 ///     </para>
 ///     <para>
 ///         This type is typically used by database providers (and other extensions). It is generally
 ///         not used in application code.
 ///     </para>
 /// </summary>
-public class PathSegment
+public struct PathSegment
 {
     /// <summary>
-    ///     Creates a new instance of the <see cref="PathSegment" /> class.
+    ///     Creates a new <see cref="PathSegment" /> struct representing JSON property access.
     /// </summary>
-    /// <param name="key">A key which is being accessed in the JSON.</param>
-    public PathSegment(string key)
+    /// <param name="propertyName">A name of JSON property which is being accessed.</param>
+    public PathSegment(string propertyName)
     {
-        Key = key;
+        PropertyName = propertyName;
+        ArrayIndex = null;
     }
 
     /// <summary>
-    ///     The key which is being accessed in the JSON.
+    ///     Creates a new <see cref="PathSegment" /> struct representing JSON array element access.
     /// </summary>
-    public virtual string Key { get; }
+    /// <param name="arrayIndex"><see langword="abstract"/>An index of an element which is being accessed in the JSON array.</param>
+    public PathSegment(SqlExpression arrayIndex)
+    {
+        ArrayIndex = arrayIndex;
+        PropertyName = null;
+    }
+
+    /// <summary>
+    ///     The name of JSON property which is being accessed.
+    /// </summary>
+    public string? PropertyName { get; }
+
+    /// <summary>
+    ///     The index of an element which is being accessed in the JSON array.
+    /// </summary>
+    public SqlExpression? ArrayIndex { get; }
 
     /// <inheritdoc />
     public override string ToString()
-        => (Key == "$" ? "" : ".") + Key;
+        => (PropertyName == "$" ? "" : ".")
+        + (PropertyName == null ? "" : PropertyName)
+        + (ArrayIndex == null ? "" : $"[{ArrayIndex}]");
 
     /// <inheritdoc />
     public override bool Equals(object? obj)
-        => obj != null
-            && (ReferenceEquals(this, obj)
-                || obj is PathSegment pathSegment
-                && Equals(pathSegment));
+        => obj is PathSegment pathSegment && Equals(pathSegment);
 
     private bool Equals(PathSegment pathSegment)
-        => Key == pathSegment.Key;
+        => PropertyName == pathSegment.PropertyName
+            && ((ArrayIndex == null && pathSegment.ArrayIndex == null)
+                || (ArrayIndex != null && ArrayIndex.Equals(pathSegment.ArrayIndex)));
 
     /// <inheritdoc />
     public override int GetHashCode()
-        => HashCode.Combine(Key);
+        => HashCode.Combine(PropertyName, ArrayIndex);
 }

src/EFCore.Relational/Query/RelationalQueryTranslationPreprocessor.cs

@@ -37,6 +37,7 @@ public override Expression NormalizeQueryableMethod(Expression expression)
         expression = new RelationalQueryMetadataExtractingExpressionVisitor(_relationalQueryCompilationContext).Visit(expression);
         expression = base.NormalizeQueryableMethod(expression);
         expression = new TableValuedFunctionToQueryRootConvertingExpressionVisitor(QueryCompilationContext.Model).Visit(expression);
+        expression = new CollectionIndexerToElementAtNormalizingExpressionVisitor().Visit(expression);
 
         return expression;
     }

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -1641,6 +1641,66 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                     ?? methodCallExpression.Update(null!, new[] { source, methodCallExpression.Arguments[1] });
             }
 
+            if (methodCallExpression.Method.IsGenericMethod
+                && (methodCallExpression.Method.GetGenericMethodDefinition() == QueryableMethods.ElementAt
+                    || methodCallExpression.Method.GetGenericMethodDefinition() == QueryableMethods.ElementAtOrDefault))
+            {
+                source = methodCallExpression.Arguments[0];
+                var selectMethodCallExpression = default(MethodCallExpression);
+
+                if (source is MethodCallExpression sourceMethodCall
+                    && sourceMethodCall.Method.IsGenericMethod
+                    && sourceMethodCall.Method.GetGenericMethodDefinition() == QueryableMethods.Select)
+                {
+                    selectMethodCallExpression = sourceMethodCall;
+                    source = sourceMethodCall.Arguments[0];
+                }
+
+                if (source is MethodCallExpression asQueryableMethodCall
+                    && asQueryableMethodCall.Method.IsGenericMethod
+                    && asQueryableMethodCall.Method.GetGenericMethodDefinition() == QueryableMethods.AsQueryable)
+                {
+                    source = asQueryableMethodCall.Arguments[0];
+                }
+
+                source = Visit(source);
+
+                if (source is JsonQueryExpression jsonQueryExpression)
+                {
+                    var collectionIndexExpression = _sqlTranslator.Translate(methodCallExpression.Arguments[1]!);
+                    if (collectionIndexExpression == null)
+                    {
+                        throw new InvalidOperationException(
+                            _sqlTranslator.TranslationErrorDetails != null
+                            ? CoreStrings.TranslationFailedWithDetails(
+                                methodCallExpression.Arguments[1],
+                                _sqlTranslator.TranslationErrorDetails)
+                            : CoreStrings.TranslationFailed(methodCallExpression.Arguments[1]));
+                    }
+
+                    collectionIndexExpression = _sqlExpressionFactory.ApplyDefaultTypeMapping(collectionIndexExpression);
+                    var newJsonQuery = jsonQueryExpression.BindCollectionElement(collectionIndexExpression!);
+
+                    var entityShaper = new RelationalEntityShaperExpression(
+                        jsonQueryExpression.EntityType,
+                        newJsonQuery,
+                        nullable: true);
+
+                    if (selectMethodCallExpression != null)
+                    {
+                        var selectorLambda = selectMethodCallExpression.Arguments[1].UnwrapLambdaFromQuote();
+                        var replaced = new ReplacingExpressionVisitor(new[] { selectorLambda.Parameters[0] }, new[] { entityShaper })
+                            .Visit(selectorLambda.Body);
+
+                        var result = Visit(replaced);
+
+                        return result;
+                    }
+
+                    return entityShaper;
+                }
+            }
+
             return base.VisitMethodCall(methodCallExpression);
         }
 

src/EFCore.Relational/Query/RelationalShapedQueryCompilingExpressionVisitor.ShaperProcessingExpressionVisitor.cs

@@ -416,9 +416,14 @@ protected override Expression VisitExtension(Expression extensionExpression)
                 {
                     if (!_variableShaperMapping.TryGetValue(entityShaperExpression.ValueBufferExpression, out var accessor))
                     {
-                        if (GetProjectionIndex(projectionBindingExpression) is ValueTuple<int, List<(IProperty, int)>, string[]>
+                        if (GetProjectionIndex(projectionBindingExpression) is ValueTuple<int, List<(IProperty, int)>, string[], int>
                             jsonProjectionIndex)
                         {
+                            if (_isTracking && jsonProjectionIndex.Item4 > 0)
+                            {
+                                throw new InvalidOperationException(RelationalStrings.ProjectingJsonCollectionElementRequiresNoTracking);
+                            }
+
                             // json entity at the root
                             var (jsonElementParameter, keyValuesParameter) = JsonShapingPreProcess(
                                 jsonProjectionIndex,
@@ -510,8 +515,13 @@ protected override Expression VisitExtension(Expression extensionExpression)
                 case CollectionResultExpression collectionResultExpression
                     when collectionResultExpression.Navigation is INavigation navigation
                     && GetProjectionIndex(collectionResultExpression.ProjectionBindingExpression)
-                        is ValueTuple<int, List<(IProperty, int)>, string[]> jsonProjectionIndex:
+                        is ValueTuple<int, List<(IProperty, int)>, string[], int> jsonProjectionIndex:
                 {
+                    if (_isTracking && jsonProjectionIndex.Item4 > 0)
+                    {
+                        throw new InvalidOperationException(RelationalStrings.ProjectingJsonCollectionElementRequiresNoTracking);
+                    }
+
                     // json entity collection at the root
                     var (jsonElementParameter, keyValuesParameter) = JsonShapingPreProcess(
                         jsonProjectionIndex,
@@ -781,7 +791,7 @@ protected override Expression VisitExtension(Expression extensionExpression)
 
                         // json include case
                         if (projectionBindingExpression != null
-                            && GetProjectionIndex(projectionBindingExpression) is ValueTuple<int, List<(IProperty, int)>, string[]>
+                            && GetProjectionIndex(projectionBindingExpression) is ValueTuple<int, List<(IProperty, int)>, string[], int>
                                 jsonProjectionIndex)
                         {
                             var (jsonElementParameter, keyValuesParameter) = JsonShapingPreProcess(
@@ -1236,16 +1246,17 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
         }
 
         private (ParameterExpression, ParameterExpression) JsonShapingPreProcess(
-            ValueTuple<int, List<(IProperty, int)>, string[]> projectionIndex,
+            ValueTuple<int, List<(IProperty, int)>, string[], int> projectionIndex,
             IEntityType entityType,
             bool isCollection)
         {
             var jsonColumnProjectionIndex = projectionIndex.Item1;
             var keyInfo = projectionIndex.Item2;
             var additionalPath = projectionIndex.Item3;
+            var specifiedCollectionIndexesCount = projectionIndex.Item4;
 
             var keyValuesParameter = Expression.Parameter(typeof(object[]));
-            var keyValues = new Expression[keyInfo.Count];
+            var keyValues = new Expression[keyInfo.Count + specifiedCollectionIndexesCount];
 
             for (var i = 0; i < keyInfo.Count; i++)
             {
@@ -1262,6 +1273,12 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                     typeof(object));
             }
 
+            for (var i = 0; i < specifiedCollectionIndexesCount; i++)
+            {
+                keyValues[keyInfo.Count + i] = Expression.Convert(
+                    Expression.Constant(0), typeof(object));
+            }
+
             var keyValuesInitialize = Expression.NewArrayInit(typeof(object), keyValues);
             var keyValuesAssignment = Expression.Assign(keyValuesParameter, keyValuesInitialize);