Make indexing work over complex JSON collections

Part of #36296

Author: roji

src/EFCore.Relational/Query/RelationalQueryableMethodTranslatingExpressionVisitor.cs

@@ -1551,6 +1551,7 @@ private sealed class SharedTypeEntityExpandingExpressionVisitor(
     {
         private readonly SqlAliasManager _sqlAliasManager = queryableTranslator._sqlAliasManager;
         private SelectExpression _selectExpression = null!;
+        private bool _bindComplexProperties;
 
         public Expression Expand(SelectExpression selectExpression, Expression lambdaBody)
         {
@@ -1564,6 +1565,7 @@ protected override Expression VisitMember(MemberExpression memberExpression)
             var innerExpression = Visit(memberExpression.Expression);
 
             return TryExpand(innerExpression, MemberIdentity.Create(memberExpression.Member))
+                ?? TryBindPropertyAccess(innerExpression, memberExpression.Member.Name)
                 ?? memberExpression.Update(innerExpression);
         }
 
@@ -1574,10 +1576,17 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                 source = Visit(source);
 
                 return TryExpand(source, MemberIdentity.Create(navigationName))
-                    ?? TryBindPrimitiveCollection(source, navigationName)
-                    ?? methodCallExpression.Update(null!, new[] { source, methodCallExpression.Arguments[1] });
+                    ?? TryBindPropertyAccess(source, navigationName)
+                    ?? methodCallExpression.Update(null!, [source, methodCallExpression.Arguments[1]]);
             }
 
+            // The following is a hack.
+            // In preprocessing, SubqueryMemberPushdownExpressionVisitor rewrites ElementAt(0).Int to Select(s => s.Int).ElementAt(0),
+            // as that's apparently needed for NavigationExpandingExpressionVisitor to work correctly;
+            // see https://github.com/dotnet/efcore/issues/30386#issuecomment-1452643142 for more context.
+            // Here, we pattern match the resulting Select(s => s.Int).ElementAt(0) over JsonQueryExpression to properly identify the
+            // indexing operation and translate it.
+            // Note that this happens for both owned and complex JSON-mapped types.
             if (methodCallExpression.Method.IsGenericMethod
                 && (methodCallExpression.Method.GetGenericMethodDefinition() == QueryableMethods.ElementAt
                     || methodCallExpression.Method.GetGenericMethodDefinition() == QueryableMethods.ElementAtOrDefault))
@@ -1600,7 +1609,16 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                     source = maybeAsQueryableMethodCall.Arguments[0];
                 }
 
+                // This special visitor normally only handles owned access, and not complex properties; all complex property access
+                // is handled in RelationalSqlTranslatingExpressionVisitor as usual.
+                // However, indexing over complex JSON-mapped collection is mangled in preprocessing just like it is for owned
+                // entities (see comment above). So we very specifically handle complex properties as part of the pattern matching
+                // hack here, to make sure indexing works properly.
+                // #36335 tracks removing the preprocessing mangling, at which point this should no longer be needed
+                var parentBindComplexProperties = _bindComplexProperties;
+                _bindComplexProperties = true;
                 source = Visit(source);
+                _bindComplexProperties = parentBindComplexProperties;
 
                 if (source is JsonQueryExpression jsonQueryExpression)
                 {
@@ -1666,15 +1684,15 @@ static Expression PrepareFailedTranslationResult(
                 var result = source;
                 if (asQueryable != null)
                 {
-                    result = asQueryable.Update(null, new[] { result });
+                    result = asQueryable.Update(null, [result]);
                 }
 
                 if (select != null)
                 {
-                    result = select.Update(null, new[] { result, select.Arguments[1] });
+                    result = select.Update(null, [result, select.Arguments[1]]);
                 }
 
-                return elementAt.Update(null, new[] { result, elementAt.Arguments[1] });
+                return elementAt.Update(null, [result, elementAt.Arguments[1]]);
             }
 
             static bool IsValidSelectorForJsonArrayElementAccess(Expression expression, JsonQueryExpression baselineJsonQuery)
@@ -1907,7 +1925,7 @@ static TableExpressionBase FindRootTableExpressionForColumn(SelectExpression sel
             }
         }
 
-        private Expression? TryBindPrimitiveCollection(Expression? source, string memberName)
+        private Expression? TryBindPropertyAccess(Expression? source, string memberName)
         {
             while (source is IncludeExpression includeExpression)
             {
@@ -1937,12 +1955,27 @@ static TableExpressionBase FindRootTableExpressionForColumn(SelectExpression sel
             }
 
             var property = type.FindProperty(memberName);
-            if (property?.IsPrimitiveCollection != true)
+            if (property?.IsPrimitiveCollection is true)
             {
-                return null;
+                return source.CreateEFPropertyExpression(property);
+            }
+
+            // See comments on indexing-related hacks in VisitMethodCall above
+            if (_bindComplexProperties
+                && type.FindComplexProperty(memberName) is IComplexProperty { IsCollection: true } complexProperty)
+            {
+                Check.DebugAssert(complexProperty.ComplexType.IsMappedToJson());
+
+                if (queryableTranslator._sqlTranslator.TryBindMember(
+                    queryableTranslator._sqlTranslator.Visit(source), MemberIdentity.Create(memberName),
+                    out var translatedExpression, out _)
+                    && translatedExpression is CollectionResultExpression { QueryExpression: JsonQueryExpression jsonQuery })
+                {
+                    return jsonQuery;
+                }
             }
 
-            return source.CreateEFPropertyExpression(property);
+            return null;
         }
 
         private sealed class AnnotationApplyingExpressionVisitor(IReadOnlyList<IAnnotation> annotations) : ExpressionVisitor

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

@@ -120,45 +120,39 @@ public override async Task Distinct_over_projected_filtered_nested_collection()
 
     public override async Task Index_constant()
     {
-        // Complex collection indexing currently fails because SubqueryMemberPushdownExpressionVisitor moves the Int member access to before the
-        // ElementAt (making a Select()), this interferes with our translation. See #36335.
-        await Assert.ThrowsAsync<EqualException>(() => base.Index_constant());
+        await base.Index_constant();
 
         AssertSql(
             """
 SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
 FROM [RootEntity] AS [r]
-WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[0]') AS int) = 8
+WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[0].Int') AS int) = 8
 """);
     }
 
     public override async Task Index_parameter()
     {
-        // Complex collection indexing currently fails because SubqueryMemberPushdownExpressionVisitor moves the Int member access to before the
-        // ElementAt (making a Select()), this interferes with our translation. See #36335.
-        await Assert.ThrowsAsync<EqualException>(() => base.Index_parameter());
+        await base.Index_parameter();
 
         AssertSql(
             """
 @i='?' (DbType = Int32)
 
 SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
 FROM [RootEntity] AS [r]
-WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[' + CAST(@i AS nvarchar(max)) + ']') AS int) = 8
+WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[' + CAST(@i AS nvarchar(max)) + '].Int') AS int) = 8
 """);
     }
 
     public override async Task Index_column()
     {
-        // Complex collection indexing currently fails because SubqueryMemberPushdownExpressionVisitor moves the Int member access to before the
-        // ElementAt (making a Select()), this interferes with our translation. See #36335.
-        await Assert.ThrowsAsync<EqualException>(() => base.Index_column());
+        await base.Index_column();
 
         AssertSql(
             """
 SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
 FROM [RootEntity] AS [r]
-WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[' + CAST([r].[Id] - 1 AS nvarchar(max)) + ']') AS int) = 8
+WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[' + CAST([r].[Id] - 1 AS nvarchar(max)) + '].Int') AS int) = 8
 """);
     }
 
@@ -170,7 +164,7 @@ public override async Task Index_out_of_bounds()
             """
 SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
 FROM [RootEntity] AS [r]
-WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[9999]') AS int) = 8
+WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[9999].Int') AS int) = 8
 """);
     }
 

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

@@ -105,6 +105,8 @@ public override async Task Distinct_over_projected_filtered_nested_collection()
 
     #endregion Distinct
 
+    #region Index
+
     public override async Task Index_constant()
     {
         await base.Index_constant();
@@ -133,6 +135,8 @@ public override async Task Index_out_of_bounds()
         AssertSql();
     }
 
+    #endregion Index
+
     public override async Task Select_within_Select_within_Select_with_aggregates()
     {
         await base.Select_within_Select_within_Select_with_aggregates();

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

@@ -40,18 +40,6 @@ [String] nvarchar(max) '$.String'
 """);
     }
 
-    public override async Task Index_constant()
-    {
-        await base.Index_constant();
-
-        AssertSql(
-            """
-SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
-FROM [RootEntity] AS [r]
-WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[0].Int') AS int) = 8
-""");
-    }
-
     public override async Task OrderBy_ElementAt()
     {
         await base.OrderBy_ElementAt();
@@ -146,6 +134,18 @@ public override async Task Distinct_over_projected_filtered_nested_collection()
 
     #region Index
 
+    public override async Task Index_constant()
+    {
+        await base.Index_constant();
+
+        AssertSql(
+            """
+SELECT [r].[Id], [r].[Name], [r].[OptionalRelated], [r].[RelatedCollection], [r].[RequiredRelated]
+FROM [RootEntity] AS [r]
+WHERE CAST(JSON_VALUE([r].[RelatedCollection], '$[0].Int') AS int) = 8
+""");
+    }
+
     public override async Task Index_parameter()
     {
         await base.Index_parameter();

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

@@ -101,7 +101,6 @@ ORDER BY [r0].[Id]
 """);
     }
 
-
     #region Distinct
 
     public override async Task Distinct()
@@ -181,6 +180,8 @@ public override async Task Distinct_over_projected_filtered_nested_collection()
 
     #endregion Distinct
 
+    #region Index
+
     public override async Task Index_constant()
     {
         await base.Index_constant();
@@ -315,6 +316,8 @@ ORDER BY (SELECT 1)
 """);
     }
 
+    #endregion Index
+
     public override async Task Select_within_Select_within_Select_with_aggregates()
     {
         await base.Select_within_Select_within_Select_with_aggregates();

test/EFCore.Sqlite.FunctionalTests/Query/Relationships/ComplexJson/ComplexJsonCollectionSqliteTest.cs

@@ -1,14 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using Microsoft.Data.Sqlite;
-
 namespace Microsoft.EntityFrameworkCore.Query.Relationships.ComplexJson;
 
 public class ComplexJsonCollectionSqliteTest(ComplexJsonSqliteFixture fixture, ITestOutputHelper testOutputHelper)
-    : ComplexJsonCollectionRelationalTestBase<ComplexJsonSqliteFixture>(fixture, testOutputHelper)
-{
-    // TODO: #36296
-    public override Task Index_column()
-        => Assert.ThrowsAsync<SqliteException>(() => base.Index_column());
-}
+    : ComplexJsonCollectionRelationalTestBase<ComplexJsonSqliteFixture>(fixture, testOutputHelper);

test/EFCore.Sqlite.FunctionalTests/Query/Relationships/ComplexJson/ComplexJsonMiscellaneousSqliteTest.cs

@@ -4,6 +4,4 @@
 namespace Microsoft.EntityFrameworkCore.Query.Relationships.ComplexJson;
 
 public class ComplexJsonMiscellaneousSqliteTest(ComplexJsonSqliteFixture fixture, ITestOutputHelper testOutputHelper)
-    : ComplexJsonMiscellaneousRelationalTestBase<ComplexJsonSqliteFixture>(fixture, testOutputHelper)
-{
-}
+    : ComplexJsonMiscellaneousRelationalTestBase<ComplexJsonSqliteFixture>(fixture, testOutputHelper);

test/EFCore.Sqlite.FunctionalTests/Query/Relationships/ComplexJson/ComplexJsonSetOperationsSqliteTest.cs

@@ -4,7 +4,4 @@
 namespace Microsoft.EntityFrameworkCore.Query.Relationships.ComplexJson;
 
 public class ComplexJsonSetOperationsSqliteTest(ComplexJsonSqliteFixture fixture, ITestOutputHelper testOutputHelper)
-    : ComplexJsonSetOperationsRelationalTestBase<ComplexJsonSqliteFixture>(fixture, testOutputHelper)
-{
-
-}
+    : ComplexJsonSetOperationsRelationalTestBase<ComplexJsonSqliteFixture>(fixture, testOutputHelper);

test/EFCore.Sqlite.FunctionalTests/Query/Relationships/ComplexTableSplitting/ComplexTableSplittingMiscellaneousSqliteTest.cs

@@ -6,6 +6,4 @@ namespace Microsoft.EntityFrameworkCore.Query.Relationships.ComplexTableSplittin
 public class ComplexTableSplittingMiscellaneousSqliteTest(
     ComplexTableSplittingSqliteFixture fixture,
     ITestOutputHelper testOutputHelper)
-    : ComplexTableSplittingMiscellaneousRelationalTestBase<ComplexTableSplittingSqliteFixture>(fixture, testOutputHelper)
-{
-}
+    : ComplexTableSplittingMiscellaneousRelationalTestBase<ComplexTableSplittingSqliteFixture>(fixture, testOutputHelper);

test/EFCore.Sqlite.FunctionalTests/Query/Relationships/Navigations/NavigationsCollectionSqliteTest.cs

@@ -4,6 +4,4 @@
 namespace Microsoft.EntityFrameworkCore.Query.Relationships.Navigations;
 
 public class NavigationsCollectionSqliteTest(NavigationsSqliteFixture fixture, ITestOutputHelper testOutputHelper)
-    : NavigationsCollectionRelationalTestBase<NavigationsSqliteFixture>(fixture, testOutputHelper)
-{
-}
+    : NavigationsCollectionRelationalTestBase<NavigationsSqliteFixture>(fixture, testOutputHelper);