Fix indexing on nested complex JSON collections

Fixes #37016

Author: roji

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -1945,12 +1945,19 @@ static TableExpressionBase FindRootTableExpressionForColumn(SelectExpression sel
             }
 
             source = source.UnwrapTypeConversion(out var convertedType);
-            if (source is not StructuralTypeShaperExpression shaper)
+
+            var type = source switch
+            {
+                StructuralTypeShaperExpression shaper => shaper.StructuralType,
+                JsonQueryExpression jsonQuery => jsonQuery.StructuralType,
+                _ => null
+            };
+
+            if (type is null)
             {
                 return null;
             }
 
-            var type = shaper.StructuralType;
             if (convertedType != null)
             {
                 Check.DebugAssert(
@@ -1969,22 +1976,50 @@ static TableExpressionBase FindRootTableExpressionForColumn(SelectExpression sel
             var property = type.FindProperty(memberName);
             if (property?.IsPrimitiveCollection is true)
             {
-                return source.CreateEFPropertyExpression(property);
+                return source!.CreateEFPropertyExpression(property);
             }
 
             // See comments on indexing-related hacks in VisitMethodCall above
-            if (_bindComplexProperties
-                && type.FindComplexProperty(memberName) is { IsCollection: true } complexProperty)
+            if (_bindComplexProperties && type.FindComplexProperty(memberName) is IComplexProperty complexProperty)
             {
-                Check.DebugAssert(complexProperty.ComplexType.IsMappedToJson());
+                Expression? translatedExpression;
+
+                if (source is JsonQueryExpression jsonSource)
+                {
+                    translatedExpression = jsonSource.BindStructuralProperty(complexProperty);
+                }
+                else if (!queryableTranslator._sqlTranslator.TryBindMember(
+                    queryableTranslator._sqlTranslator.Visit(source), MemberIdentity.Create(memberName),
+                    out translatedExpression, out _))
+                {
+                    return null;
+                }
 
-                if (queryableTranslator._sqlTranslator.TryBindMember(
-                        queryableTranslator._sqlTranslator.Visit(source), MemberIdentity.Create(memberName),
-                        out var translatedExpression, out _)
-                    && translatedExpression is CollectionResultExpression { QueryExpression: JsonQueryExpression jsonQuery })
+                // Hack: when returning a StructuralTypeShaperExpression, _sqlTranslator returns it wrapped by a
+                // StructuralTypeReferenceExpression, which is supposed to be a private wrapper only within the SQL translator.
+                // Call TranslateProjection to unwrap it (need to look into getting rid StructuralTypeReferenceExpression altogether).
+                if (translatedExpression is not JsonQueryExpression and not CollectionResultExpression)
                 {
-                    return jsonQuery;
+                    if (queryableTranslator._sqlTranslator.TranslateProjection(translatedExpression) is { } unwrappedTarget)
+                    {
+                        translatedExpression = unwrappedTarget;
+                    }
+                    else
+                    {
+                        return null;
+                    }
                 }
+
+                return complexProperty switch
+                {
+                    { IsCollection: false } when translatedExpression is StructuralTypeShaperExpression { ValueBufferExpression: JsonQueryExpression jsonQuery }
+                        => jsonQuery,
+                    { IsCollection: true } when translatedExpression is CollectionResultExpression { QueryExpression: JsonQueryExpression jsonQuery }
+                        => jsonQuery,
+                    { IsCollection: true } when translatedExpression is JsonQueryExpression jsonQuery
+                        => jsonQuery,
+                    _ => null
+                };
             }
 
             return null;

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

@@ -377,7 +377,7 @@ IComplexType complexType
                 } selectExpression
                     when TranslateExpression(index) is { } translatedIndex
                     && _sqlServerSingletonOptions.SupportsJsonFunctions
-                    && TryTranslate(selectExpression, valuesParameter, translatedIndex, out var result):
+                    && TryTranslate(selectExpression, valuesParameter, path: null, translatedIndex, out var result):
                     return result;
 
                 // Index on JSON array
@@ -406,7 +406,7 @@ when TranslateExpression(index) is { } translatedIndex
                 } selectExpression
                     when orderingTableAlias == openJsonExpression.Alias
                     && TranslateExpression(index) is { } translatedIndex
-                    && TryTranslate(selectExpression, jsonArrayColumn, translatedIndex, out var result):
+                    && TryTranslate(selectExpression, jsonArrayColumn, openJsonExpression.Path, translatedIndex, out var result):
                     return result;
             }
         }
@@ -415,7 +415,8 @@ when TranslateExpression(index) is { } translatedIndex
 
         bool TryTranslate(
             SelectExpression selectExpression,
-            SqlExpression jsonArrayColumn,
+            SqlExpression jsonColumn,
+            IReadOnlyList<PathSegment>? path,
             SqlExpression translatedIndex,
             [NotNullWhen(true)] out ShapedQueryExpression? result)
         {
@@ -441,16 +442,22 @@ bool TryTranslate(
                 return false;
             }
 
-            // If the inner expression happens to itself be a JsonScalarExpression, simply append the two paths to avoid creating
+            // If the inner expression happens to itself be a JsonScalarExpression, simply append the paths to avoid creating
             // JSON_VALUE within JSON_VALUE.
-            var (json, path) = jsonArrayColumn is JsonScalarExpression innerJsonScalarExpression
-                ? (innerJsonScalarExpression.Json,
-                    innerJsonScalarExpression.Path.Append(new(translatedIndex)).ToArray())
-                : (jsonArrayColumn, [new(translatedIndex)]);
+            var (json, newPath) = jsonColumn is JsonScalarExpression innerJsonScalarExpression
+                ? (innerJsonScalarExpression.Json, new List<PathSegment>(innerJsonScalarExpression.Path))
+                : (jsonColumn, []);
+
+            if (path is not null)
+            {
+                newPath.AddRange(path);
+            }
+
+            newPath.Add(new(translatedIndex));
 
             var translation = new JsonScalarExpression(
                 json,
-                path,
+                newPath,
                 projection.Type,
                 projection.TypeMapping,
                 projectionColumn.IsNullable);

test/EFCore.Cosmos.FunctionalTests/Query/Associations/OwnedNavigations/OwnedNavigationsCollectionCosmosTest.cs

@@ -126,6 +126,18 @@ FROM root c
 """);
     }
 
+    public override async Task Index_on_nested_collection()
+    {
+        await base.Index_on_nested_collection();
+
+        AssertSql(
+            """
+SELECT VALUE c
+FROM root c
+WHERE (c["RequiredAssociate"]["NestedCollection"][0]["Int"] = 8)
+""");
+    }
+
     #endregion Index
 
     #region GroupBy

test/EFCore.Specification.Tests/Query/Associations/AssociationsCollectionTestBase.cs

@@ -75,6 +75,14 @@ public virtual Task Index_column()
             ss => ss.Set<RootEntity>().Where(e => e.AssociateCollection[e.Id - 1].Int == 8),
             ss => ss.Set<RootEntity>().Where(e => e.AssociateCollection.Count > e.Id - 1 && e.AssociateCollection[e.Id - 1].Int == 8)));
 
+    [ConditionalFact]
+    public virtual Task Index_on_nested_collection()
+        => AssertOrderedCollectionQuery(() => AssertQuery(
+            ss => ss.Set<RootEntity>().Where(e => e.RequiredAssociate.NestedCollection[0].Int == 8),
+            ss => ss.Set<RootEntity>().Where(
+                e => e.RequiredAssociate.NestedCollection.Count > 0
+                    && e.RequiredAssociate.NestedCollection[0].Int == 8)));
+
     [ConditionalFact]
     public virtual Task Index_out_of_bounds()
         => AssertOrderedCollectionQuery(() => AssertQuery(

test/EFCore.SqlServer.FunctionalTests/Query/Associations/ComplexJson/ComplexJsonCollectionSqlServerTest.cs

@@ -244,6 +244,30 @@ WHERE CAST(JSON_VALUE([r].[AssociateCollection], '$[' + CAST([r].[Id] - 1 AS nva
         }
     }
 
+    public override async Task Index_on_nested_collection()
+    {
+        await base.Index_on_nested_collection();
+
+        if (Fixture.UsingJsonType)
+        {
+            AssertSql(
+                """
+SELECT [r].[Id], [r].[Name], [r].[AssociateCollection], [r].[OptionalAssociate], [r].[RequiredAssociate]
+FROM [RootEntity] AS [r]
+WHERE JSON_VALUE([r].[RequiredAssociate], '$.NestedCollection[0].Int' RETURNING int) = 8
+""");
+        }
+        else
+        {
+            AssertSql(
+                """
+SELECT [r].[Id], [r].[Name], [r].[AssociateCollection], [r].[OptionalAssociate], [r].[RequiredAssociate]
+FROM [RootEntity] AS [r]
+WHERE CAST(JSON_VALUE([r].[RequiredAssociate], '$.NestedCollection[0].Int') AS int) = 8
+""");
+        }
+    }
+
     public override async Task Index_out_of_bounds()
     {
         await base.Index_out_of_bounds();

test/EFCore.SqlServer.FunctionalTests/Query/Associations/Navigations/NavigationsCollectionSqlServerTest.cs

@@ -199,6 +199,13 @@ public override async Task Index_column()
         AssertSql();
     }
 
+    public override async Task Index_on_nested_collection()
+    {
+        await base.Index_on_nested_collection();
+
+        AssertSql();
+    }
+
     public override async Task Index_out_of_bounds()
     {
         await base.Index_out_of_bounds();

test/EFCore.SqlServer.FunctionalTests/Query/Associations/OwnedJson/OwnedJsonCollectionSqlServerTest.cs

@@ -317,6 +317,30 @@ WHERE CAST(JSON_VALUE([r].[AssociateCollection], '$[' + CAST([r].[Id] - 1 AS nva
         }
     }
 
+    public override async Task Index_on_nested_collection()
+    {
+        await base.Index_on_nested_collection();
+
+        if (Fixture.UsingJsonType)
+        {
+            AssertSql(
+                """
+SELECT [r].[Id], [r].[Name], [r].[AssociateCollection], [r].[OptionalAssociate], [r].[RequiredAssociate]
+FROM [RootEntity] AS [r]
+WHERE JSON_VALUE([r].[RequiredAssociate], '$.NestedCollection[0].Int' RETURNING int) = 8
+""");
+        }
+        else
+        {
+            AssertSql(
+                """
+SELECT [r].[Id], [r].[Name], [r].[AssociateCollection], [r].[OptionalAssociate], [r].[RequiredAssociate]
+FROM [RootEntity] AS [r]
+WHERE CAST(JSON_VALUE([r].[RequiredAssociate], '$.NestedCollection[0].Int') AS int) = 8
+""");
+        }
+    }
+
     public override async Task Index_out_of_bounds()
     {
         await base.Index_out_of_bounds();

test/EFCore.SqlServer.FunctionalTests/Query/Associations/OwnedNavigations/OwnedNavigationsCollectionSqlServerTest.cs

@@ -202,6 +202,13 @@ public override async Task Index_column()
         AssertSql();
     }
 
+    public override async Task Index_on_nested_collection()
+    {
+        await base.Index_on_nested_collection();
+
+        AssertSql();
+    }
+
     public override async Task Index_out_of_bounds()
     {
         await base.Index_out_of_bounds();