Cosmos: translate Contains over subquery

Part of dotnet#25765
roji · Jun 5, 2024 · 7b01336 · 7b01336
1 parent e62a2cb
commit 7b01336
Show file tree

Hide file tree

Showing 7 changed files with 156 additions and 30 deletions.
diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryableMethodTranslatingExpressionVisitor.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryableMethodTranslatingExpressionVisitor.cs
@@ -431,8 +431,13 @@ protected override ShapedQueryExpression TranslateCast(ShapedQueryExpression sou
             return source.UpdateQueryExpression(new SelectExpression(simplifiedTranslation));
         }
 
-        // TODO: Translation to IN, with scalars and with subquery
-        return null;
+        // Translate to EXISTS
+        var anyLambdaParameter = Expression.Parameter(item.Type, "p");
+        var anyLambda = Expression.Lambda(
+            Microsoft.EntityFrameworkCore.Infrastructure.ExpressionExtensions.CreateEqualsExpression(anyLambdaParameter, item),
+            anyLambdaParameter);
+
+        return TranslateAny(source, anyLambda);
     }
 
     /// <summary>

diff --git a/test/EFCore.Cosmos.FunctionalTests/Query/NorthwindAggregateOperatorsQueryCosmosTest.cs b/test/EFCore.Cosmos.FunctionalTests/Query/NorthwindAggregateOperatorsQueryCosmosTest.cs
@@ -1645,25 +1645,52 @@ FROM root c
 """);
             });
 
-    public override async Task Contains_with_local_enumerable_inline(bool async)
-    {
-        // Issue #31776
-        await Assert.ThrowsAsync<InvalidOperationException>(
-            async () =>
-                await base.Contains_with_local_enumerable_inline(async));
+    public override Task Contains_with_local_enumerable_inline(bool async)
+        => Fixture.NoSyncTest(
+            async, async a =>
+            {
+                await base.Contains_with_local_enumerable_inline(a);
 
-        AssertSql();
-    }
+                AssertSql(
+                    """
+SELECT c
+FROM root c
+WHERE ((c["Discriminator"] = "Customer") AND EXISTS (
+    SELECT 1
+    FROM i IN (SELECT VALUE ["ABCDE", "ALFKI"])
+    WHERE ((i != null) AND (i = c["CustomerID"]))))
+""");
+            });
 
-    public override async Task Contains_with_local_enumerable_inline_closure_mix(bool async)
-    {
-        // Issue #31776
-        await Assert.ThrowsAsync<InvalidOperationException>(
-            async () =>
-                await base.Contains_with_local_enumerable_inline_closure_mix(async));
+    public override Task Contains_with_local_enumerable_inline_closure_mix(bool async)
+        => Fixture.NoSyncTest(
+            async, async a =>
+            {
+                await base.Contains_with_local_enumerable_inline_closure_mix(a);
 
-        AssertSql();
-    }
+                AssertSql(
+                    """
+@__p_0='["ABCDE","ALFKI"]'
+
+SELECT c
+FROM root c
+WHERE ((c["Discriminator"] = "Customer") AND EXISTS (
+    SELECT 1
+    FROM i IN (SELECT VALUE @__p_0)
+    WHERE ((i != null) AND (i = c["CustomerID"]))))
+""",
+                    //
+                    """
+@__p_0='["ABCDE","ANATR"]'
+
+SELECT c
+FROM root c
+WHERE ((c["Discriminator"] = "Customer") AND EXISTS (
+    SELECT 1
+    FROM i IN (SELECT VALUE @__p_0)
+    WHERE ((i != null) AND (i = c["CustomerID"]))))
+""");
+            });
 
     public override Task Contains_with_local_ordered_enumerable_closure(bool async)
         => Fixture.NoSyncTest(
@@ -1974,10 +2001,24 @@ FROM root c
 
     public override async Task Contains_top_level(bool async)
     {
-        // Contains over subquery. Issue #17246.
-        await AssertTranslationFailed(() => base.Contains_top_level(async));
+        // Always throws for sync.
+        if (async)
+        {
+            // Top-level Any(), see #33854.
+            var exception = await Assert.ThrowsAsync<CosmosException>(() => base.Contains_top_level(async));
 
-        AssertSql();
+            Assert.Contains("Identifier 'root' could not be resolved.", exception.Message);
+
+            AssertSql(
+                """
+@__p_0='ALFKI'
+
+SELECT EXISTS (
+    SELECT 1
+    FROM root c
+    WHERE ((c["Discriminator"] = "Customer") AND (c["CustomerID"] = @__p_0))) AS c
+""");
+        }
     }
 
     public override async Task Contains_with_local_tuple_array_closure(bool async)
@@ -2204,10 +2245,25 @@ FROM root c
 
     public override async Task Contains_over_entityType_with_null_should_rewrite_to_false(bool async)
     {
-        // Contains over subquery. Issue #17246.
-        await AssertTranslationFailed(() => base.Contains_over_entityType_with_null_should_rewrite_to_false(async));
+        // Always throws for sync.
+        if (async)
+        {
+            // Top-level Any(), see #33854.
+            var exception =
+                await Assert.ThrowsAsync<CosmosException>(() => base.Contains_over_entityType_with_null_should_rewrite_to_false(async));
 
-        AssertSql();
+            Assert.Contains("Identifier 'root' could not be resolved.", exception.Message);
+
+            AssertSql(
+                """
+@__entity_equality_p_0_OrderID=null
+
+SELECT EXISTS (
+    SELECT 1
+    FROM root c
+    WHERE (((c["Discriminator"] = "Order") AND (c["CustomerID"] = "VINET")) AND (c["OrderID"] = @__entity_equality_p_0_OrderID))) AS c
+""");
+        }
     }
 
     public override async Task Contains_over_entityType_with_null_in_projection(bool async)

diff --git a/test/EFCore.Cosmos.FunctionalTests/Query/PrimitiveCollectionsQueryCosmosTest.cs b/test/EFCore.Cosmos.FunctionalTests/Query/PrimitiveCollectionsQueryCosmosTest.cs
@@ -985,20 +985,43 @@ public override async Task Column_collection_Skip(bool async)
 
     public override async Task Column_collection_Take(bool async)
     {
-        // TODO: IN with subquery
-        await AssertTranslationFailed(() => base.Column_collection_Take(async));
+        // Always throws for sync.
+        if (async)
+        {
+            var exception = await Assert.ThrowsAsync<CosmosException>(() => base.Column_collection_Take(async));
 
-        AssertSql();
+            Assert.Contains("'OFFSET LIMIT' clause is not supported in subqueries.", exception.Message);
+        }
     }
 
     public override async Task Column_collection_Skip_Take(bool async)
     {
-        // TODO: Count after Distinct requires subquery pushdown
-        await AssertTranslationFailed(() => base.Column_collection_Skip_Take(async));
+        // Always throws for sync.
+        if (async)
+        {
+            var exception = await Assert.ThrowsAsync<CosmosException>(() => base.Column_collection_Skip_Take(async));
 
-        AssertSql();
+            Assert.Contains("'OFFSET LIMIT' clause is not supported in subqueries.", exception.Message);
+        }
     }
 
+    public override Task Column_collection_Contains_over_subquery(bool async)
+        => CosmosTestHelpers.Instance.NoSyncTest(
+            async, async a =>
+            {
+                await base.Column_collection_Contains_over_subquery(a);
+
+                AssertSql(
+                    """
+SELECT c
+FROM root c
+WHERE ((c["Discriminator"] = "PrimitiveCollectionsEntity") AND EXISTS (
+    SELECT 1
+    FROM i IN c["Ints"]
+    WHERE ((i > 1) AND (i = 11))))
+""");
+            });
+
     public override async Task Column_collection_OrderByDescending_ElementAt(bool async)
     {
         // TODO: ElementAt over composed query (non-simple array)

diff --git a/test/EFCore.Specification.Tests/Query/PrimitiveCollectionsQueryTestBase.cs b/test/EFCore.Specification.Tests/Query/PrimitiveCollectionsQueryTestBase.cs
@@ -605,6 +605,13 @@ public virtual Task Column_collection_Skip_Take(bool async)
             async,
             ss => ss.Set<PrimitiveCollectionsEntity>().Where(c => c.Ints.Skip(1).Take(2).Contains(11)));
 
+    [ConditionalTheory]
+    [MemberData(nameof(IsAsyncData))]
+    public virtual Task Column_collection_Contains_over_subquery(bool async)
+        => AssertQuery(
+            async,
+            ss => ss.Set<PrimitiveCollectionsEntity>().Where(c => c.Ints.Where(i => i > 1).Contains(11)));
+
     [ConditionalTheory]
     [MemberData(nameof(IsAsyncData))]
     public virtual Task Column_collection_OrderByDescending_ElementAt(bool async)

diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQueryOldSqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQueryOldSqlServerTest.cs
@@ -688,6 +688,9 @@ public override Task Column_collection_Take(bool async)
     public override Task Column_collection_Skip_Take(bool async)
         => AssertCompatibilityLevelTooLow(() => base.Column_collection_Skip_Take(async));
 
+    public override Task Column_collection_Contains_over_subquery(bool async)
+        => AssertCompatibilityLevelTooLow(() => base.Column_collection_Skip_Take(async));
+
     public override Task Column_collection_OrderByDescending_ElementAt(bool async)
         => AssertTranslationFailed(() => base.Column_collection_OrderByDescending_ElementAt(async));
 

diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/PrimitiveCollectionsQuerySqlServerTest.cs
@@ -1004,6 +1004,22 @@ OFFSET 1 ROWS FETCH NEXT 2 ROWS ONLY
 """);
     }
 
+    public override async Task Column_collection_Contains_over_subquery(bool async)
+    {
+        await base.Column_collection_Contains_over_subquery(async);
+
+        AssertSql(
+            """
+SELECT [p].[Id], [p].[Bool], [p].[Bools], [p].[DateTime], [p].[DateTimes], [p].[Enum], [p].[Enums], [p].[Int], [p].[Ints], [p].[NullableInt], [p].[NullableInts], [p].[NullableString], [p].[NullableStrings], [p].[String], [p].[Strings]
+FROM [PrimitiveCollectionsEntity] AS [p]
+WHERE 11 IN (
+    SELECT [i].[value]
+    FROM OPENJSON([p].[Ints]) WITH ([value] int '$') AS [i]
+    WHERE [i].[value] > 1
+)
+""");
+    }
+
     public override async Task Column_collection_OrderByDescending_ElementAt(bool async)
     {
         await base.Column_collection_OrderByDescending_ElementAt(async);

diff --git a/test/EFCore.Sqlite.FunctionalTests/Query/PrimitiveCollectionsQuerySqliteTest.cs b/test/EFCore.Sqlite.FunctionalTests/Query/PrimitiveCollectionsQuerySqliteTest.cs
@@ -986,6 +986,22 @@ LIMIT 2 OFFSET 1
 """);
     }
 
+    public override async Task Column_collection_Contains_over_subquery(bool async)
+    {
+        await base.Column_collection_Contains_over_subquery(async);
+
+        AssertSql(
+            """
+SELECT "p"."Id", "p"."Bool", "p"."Bools", "p"."DateTime", "p"."DateTimes", "p"."Enum", "p"."Enums", "p"."Int", "p"."Ints", "p"."NullableInt", "p"."NullableInts", "p"."NullableString", "p"."NullableStrings", "p"."String", "p"."Strings"
+FROM "PrimitiveCollectionsEntity" AS "p"
+WHERE 11 IN (
+    SELECT "i"."value"
+    FROM json_each("p"."Ints") AS "i"
+    WHERE "i"."value" > 1
+)
+""");
+    }
+
     public override async Task Column_collection_OrderByDescending_ElementAt(bool async)
     {
         await base.Column_collection_OrderByDescending_ElementAt(async);