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>Query is projecting entity mapped to JSON that accesses a JSON collection element. Those 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/CollectionIndexerToElementAtConvertingExpressionVisitor.cs

@@ -0,0 +1,59 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Linq.Expressions;
+
+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 CollectionIndexerToElementAtConvertingExpressionVisitor : ExpressionVisitor
+{
+    protected override Expression VisitMethodCall(MethodCallExpression methodCallExpression)
+    {
+        if (methodCallExpression.Method.Name == "get_Item"
+            && !methodCallExpression.Method.IsStatic
+            && methodCallExpression.Method.DeclaringType != null
+            && ((methodCallExpression.Method.DeclaringType.IsGenericType
+                && methodCallExpression.Method.DeclaringType.GetGenericTypeDefinition() == typeof(List<>))
+                || methodCallExpression.Method.DeclaringType.IsArray))
+        {
+            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)
+    {
+        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,30 @@ 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)
+    {
+        Debug.Assert(
+            Path.Last().CollectionIndexExpression == null,
+            "Already accessing collection 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

@@ -17,22 +17,40 @@ namespace Microsoft.EntityFrameworkCore.Query;
 public class PathSegment
 {
     /// <summary>
-    ///     Creates a new instance of the <see cref="PathSegment" /> class.
+    ///     Creates a new instance of the <see cref="PathSegment" /> class representing JSON property access.
     /// </summary>
     /// <param name="key">A key which is being accessed in the JSON.</param>
     public PathSegment(string key)
     {
         Key = key;
+        CollectionIndexExpression = null;
+    }
+
+    /// <summary>
+    ///     Creates a new instance of the <see cref="PathSegment" /> class representing JSON collection element access.
+    /// </summary>
+    /// <param name="collectionIndexExpression">A collection index which is being accessed in the JSON.</param>
+    public PathSegment(SqlExpression collectionIndexExpression)
+    {
+        CollectionIndexExpression = collectionIndexExpression;
+        Key = null;
     }
 
     /// <summary>
     ///     The key which is being accessed in the JSON.
     /// </summary>
-    public virtual string Key { get; }
+    public virtual string? Key { get; }
+
+    /// <summary>
+    ///     The index of the collection which is being accessed in the JSON.
+    /// </summary>
+    public virtual SqlExpression? CollectionIndexExpression { get; }
 
     /// <inheritdoc />
     public override string ToString()
-        => (Key == "$" ? "" : ".") + Key;
+        => (Key == "$" ? "" : ".")
+        + (Key == null ? "" : Key)
+        + (CollectionIndexExpression == null ? "" : $"[{CollectionIndexExpression}]");
 
     /// <inheritdoc />
     public override bool Equals(object? obj)
@@ -42,9 +60,11 @@ public override bool Equals(object? obj)
                 && Equals(pathSegment));
 
     private bool Equals(PathSegment pathSegment)
-        => Key == pathSegment.Key;
+        => Key == pathSegment.Key
+            && ((CollectionIndexExpression == null && pathSegment.CollectionIndexExpression == null)
+                || (CollectionIndexExpression != null && CollectionIndexExpression.Equals(pathSegment.CollectionIndexExpression)));
 
     /// <inheritdoc />
     public override int GetHashCode()
-        => HashCode.Combine(Key);
+        => HashCode.Combine(Key, CollectionIndexExpression);
 }

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 CollectionIndexerToElementAtConvertingExpressionVisitor().Visit(expression);
 
         return expression;
     }

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -1640,6 +1640,64 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                 return TryExpand(source, MemberIdentity.Create(navigationName))
                     ?? methodCallExpression.Update(null!, new[] { source, methodCallExpression.Arguments[1] });
             }
+            else 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 sourceMethod
+                    && sourceMethod.Method.IsGenericMethod
+                    && sourceMethod.Method.GetGenericMethodDefinition() == QueryableMethods.Select)
+                {
+                    selectMethodCallExpression = sourceMethod;
+                    source = sourceMethod.Arguments[0];
+                }
+
+                if (source is MethodCallExpression asQueryableMethodCallExpression
+                    && asQueryableMethodCallExpression.Method.IsGenericMethod
+                    && asQueryableMethodCallExpression.Method.GetGenericMethodDefinition() == QueryableMethods.AsQueryable)
+                {
+                    source = asQueryableMethodCallExpression.Arguments[0];
+                }
+
+                source = Visit(source);
+
+                if (source is JsonQueryExpression jsonQueryExpression)
+                {
+                    var collectionIndexExpression = _sqlTranslator.Translate(methodCallExpression.Arguments[1]!);
+
+                    if (collectionIndexExpression == null)
+                    {
+                        return methodCallExpression.Update(null!, new[] { source, methodCallExpression.Arguments[1] });
+                    }
+
+                    collectionIndexExpression = _sqlExpressionFactory.ApplyDefaultTypeMapping(collectionIndexExpression);
+                    var newJsonQuery = jsonQueryExpression.BindCollectionElement(collectionIndexExpression!);
+
+                    var entityShaper = new RelationalEntityShaperExpression(
+                        jsonQueryExpression.EntityType,
+                        newJsonQuery,
+                        nullable: true);
+
+                    // look into select (if there was any)
+                    // strip the includes
+                    // and if there was anything (e.g. MaterializeCollectionNavigationExpression) wrap the entity shaper around it and return that
+                    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 @@ when collectionResultExpression.Navigation is INavigation navigation
 
                         // 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 @@ private Expression CreateJsonShapers(
         }
 
         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 @@ private Expression CreateJsonShapers(
                     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);
 

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

@@ -1529,15 +1529,15 @@ Expression CopyProjectionToOuter(SelectExpression innerSelectExpression, Express
 
                         remappedConstant = Constant(newDictionary);
                     }
-                    else if (constantValue is ValueTuple<int, List<ValueTuple<IProperty, int>>, string[]> tuple)
+                    else if (constantValue is ValueTuple<int, List<(IProperty, int)>, string[], int> tuple)
                     {
-                        var newList = new List<ValueTuple<IProperty, int>>();
+                        var newList = new List<(IProperty, int)>();
                         foreach (var item in tuple.Item2)
                         {
                             newList.Add((item.Item1, projectionIndexMap[item.Item2]));
                         }
 
-                        remappedConstant = Constant((projectionIndexMap[tuple.Item1], newList, tuple.Item3));
+                        remappedConstant = Constant((projectionIndexMap[tuple.Item1], newList, tuple.Item3, tuple.Item4));
                     }
                     else
                     {
@@ -1590,7 +1590,8 @@ static Dictionary<JsonQueryExpression, JsonScalarExpression> BuildJsonProjection
             {
                 var ordered = projections
                     .OrderBy(x => $"{x.JsonColumn.TableAlias}.{x.JsonColumn.Name}")
-                    .ThenBy(x => x.Path.Count);
+                    .ThenBy(x => x.Path.Count)
+                    .ThenBy(x => x.Path.Last().CollectionIndexExpression != null);
 
                 var needed = new List<JsonScalarExpression>();
                 foreach (var orderedElement in ordered)
@@ -1640,10 +1641,13 @@ ConstantExpression AddJsonProjection(JsonQueryExpression jsonQueryExpression, Js
         {
             var additionalPath = new string[0];
 
-            // this will be more tricky once we support more complicated json path options
+            // TODO: change this when implementing #29513
+            // deduplication doesn't happen currently if the additional path contains indexes
+            Debug.Assert(jsonQueryExpression.Path.Skip(jsonScalarToAdd.Path.Count).All(x => x.Key != null));
+
             additionalPath = jsonQueryExpression.Path
                 .Skip(jsonScalarToAdd.Path.Count)
-                .Select(x => x.Key)
+                .Select(x => x.Key!)
                 .ToArray();
 
             var jsonColumnIndex = AddToProjection(jsonScalarToAdd);
@@ -1656,7 +1660,9 @@ ConstantExpression AddJsonProjection(JsonQueryExpression jsonQueryExpression, Js
                 keyInfo.Add((keyProperty, AddToProjection(keyColumn)));
             }
 
-            return Constant((jsonColumnIndex, keyInfo, additionalPath));
+            var specifiedCollectionIndexesCount = jsonScalarToAdd.Path.Count(x => x.CollectionIndexExpression != null);
+
+            return Constant((jsonColumnIndex, keyInfo, additionalPath, specifiedCollectionIndexesCount));
         }
 
         static IReadOnlyList<IProperty> GetMappedKeyProperties(IKey key)
@@ -1697,7 +1703,15 @@ static bool JsonEntityContainedIn(JsonScalarExpression sourceExpression, JsonQue
                 return false;
             }
 
-            return sourcePath.SequenceEqual(targetPath.Take(sourcePath.Count));
+            if (!sourcePath.SequenceEqual(targetPath.Take(sourcePath.Count)))
+            {
+                return false;
+            }
+
+            // we can only perform deduplication if there additional path doesn't contain any collection indexes
+            // collection indexes can only be part of the source path
+            // see issue #29513
+            return targetPath.Skip(sourcePath.Count).All(x => x.CollectionIndexExpression == null);
         }
     }