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

Converting indexer into AsQueryable().ElementAt(x) so that nav expansion can understand it and inject MaterializeCollectionNavigationExpression in the right places.
Then in translation we recognize the pattern and convert it back to element access on a JsonQueryExpression.
In order to shape the results correctly, we need to add all the array indexes (that are not constants) to the projection, so that we can populate the ordinal keys correctly (and also to do de-duplication).

Fixes #28648

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

@@ -0,0 +1,39 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+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 != typeof(string)
+            && methodCallExpression.Method.DeclaringType != typeof(byte[])
+            && ((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);
+    }
+}

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)
+    {
+        var newPath = Path.Take(Path.Count - 1).ToList();
+        var lastPathSegment = Path.Last();
+        if (lastPathSegment.CollectionIndexExpression != null)
+        {
+            throw new InvalidOperationException("Already accessing collection element.");
+        }
+
+        newPath.Add(new PathSegment(lastPathSegment.Key, 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

@@ -25,14 +25,30 @@ public PathSegment(string key)
         Key = key;
     }
 
+    /// <summary>
+    ///     Creates a new instance of the <see cref="PathSegment" /> class.
+    /// </summary>
+    /// <param name="key">A key which is being accessed in the JSON.</param>
+    /// <param name="collectionIndexExpression">A collection index which is being accessed in the JSON.</param>
+    public PathSegment(string key, SqlExpression collectionIndexExpression)
+        : this(key)
+    {
+        CollectionIndexExpression = collectionIndexExpression;
+    }
+
     /// <summary>
     ///     The key which is being accessed in the JSON.
     /// </summary>
     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 + (CollectionIndexExpression == null ? "" : $"[{CollectionIndexExpression}]");
 
     /// <inheritdoc />
     public override bool Equals(object? obj)
@@ -42,9 +58,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,65 @@ 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,7 +416,7 @@ 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)
                         {
                             // json entity at the root
@@ -510,7 +510,7 @@ 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:
                 {
                     // json entity collection at the root
                     var (jsonElementParameter, keyValuesParameter) = JsonShapingPreProcess(
@@ -781,7 +781,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 +1236,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 +1263,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
                     {
@@ -1591,6 +1591,7 @@ 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.Last().CollectionIndexExpression == null);
 
                 var needed = new List<JsonScalarExpression>();
                 foreach (var orderedElement in ordered)
@@ -1656,7 +1657,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 +1700,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);
         }
     }
 

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

@@ -318,7 +318,32 @@ protected override Expression VisitJsonScalar(JsonScalarExpression jsonScalarExp
 
         Visit(jsonScalarExpression.JsonColumn);
 
-        Sql.Append($",'{string.Join("", jsonScalarExpression.Path.Select(e => e.ToString()))}')");
+        Sql.Append(",'");
+        foreach (var pathSegment in jsonScalarExpression.Path)
+        {
+            Sql.Append((pathSegment.Key == "$" ? "" : ".") + pathSegment.Key);
+            if (pathSegment.CollectionIndexExpression != null)
+            {
+                Sql.Append("[");
+
+                if (pathSegment.CollectionIndexExpression is SqlConstantExpression)
+                {
+                    Visit(pathSegment.CollectionIndexExpression);
+                }
+                else
+                {
+                    Sql.Append("' + CAST(");
+                    Visit(pathSegment.CollectionIndexExpression);
+                    Sql.Append(" AS ");
+                    Sql.Append(_typeMappingSource.GetMapping(typeof(string)).StoreType);
+                    Sql.Append(") + '");
+                }
+
+                Sql.Append("]");
+            }
+        }
+
+        Sql.Append("')");
 
         if (jsonScalarExpression.Type != typeof(JsonElement))
         {

src/EFCore/Query/Internal/NavigationExpandingExpressionVisitor.cs

@@ -619,8 +619,6 @@ when QueryableMethods.IsSumWithSelector(method):
                         // GroupJoin overloads
                         // Zip
                         // SequenceEqual overloads
-                        // ElementAt
-                        // ElementAtOrDefault
                         // SkipWhile
                         // TakeWhile
                         // DefaultIfEmpty with argument

src/EFCore/Query/Internal/QueryableMethodNormalizingExpressionVisitor.cs

@@ -125,13 +125,36 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
             visitedExpression = base.VisitMethodCall(methodCallExpression);
         }
 
-        if (visitedExpression is MethodCallExpression visitedMethodCall
-            && visitedMethodCall.Method.DeclaringType == typeof(Queryable)
-            && visitedMethodCall.Method.IsGenericMethod)
+        if (visitedExpression is MethodCallExpression visitedMethodCall)
         {
-            return TryFlattenGroupJoinSelectMany(visitedMethodCall);
+            if (visitedMethodCall.Method.DeclaringType == typeof(Queryable)
+                && visitedMethodCall.Method.IsGenericMethod)
+            {
+                return TryFlattenGroupJoinSelectMany(visitedMethodCall);
+            }
+
+            if (visitedMethodCall.Method is MethodInfo methodInfo
+                && methodInfo.Name == "get_Item"
+                && !methodInfo.IsStatic
+                && ((methodInfo.DeclaringType!.IsGenericType && methodInfo.DeclaringType.GetGenericTypeDefinition() == typeof(List<>))
+                    || methodInfo.DeclaringType.IsArray))
+            {
+                return Expression.Call(
+                    QueryableMethods.ElementAt.MakeGenericMethod(visitedMethodCall.Type),
+                    Expression.Call(
+                        QueryableMethods.AsQueryable.MakeGenericMethod(visitedMethodCall.Type),
+                        visitedMethodCall.Object!),
+                    visitedMethodCall.Arguments[0]);
+            }
         }
 
+        //if (visitedExpression is MethodCallExpression visitedMethodCall
+        //    && visitedMethodCall.Method.DeclaringType == typeof(Queryable)
+        //    && visitedMethodCall.Method.IsGenericMethod)
+        //{
+        //    return TryFlattenGroupJoinSelectMany(visitedMethodCall);
+        //}
+
         return visitedExpression;
     }