Cosmos: Implement SelectMany (#34013)

Also introduces substantial infrastructure for general joins

Closes #17312

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -123,6 +123,9 @@
   <data name="CanConnectNotSupported" xml:space="preserve">
     <value>The Cosmos database does not support 'CanConnect' or 'CanConnectAsync'.</value>
   </data>
+  <data name="ComplexProjectionInSubqueryNotSupported" xml:space="preserve">
+    <value>Complex projections in subqueries are currently unsupported.</value>
+  </data>
   <data name="ConnectionInfoMissing" xml:space="preserve">
     <value>None of connection string, CredentialToken, account key or account endpoint were specified. Specify a set of connection details.</value>
   </data>

src/EFCore.Cosmos/Query/Internal/CosmosQuerySqlGenerator.cs

@@ -283,16 +283,11 @@ protected override Expression VisitSelect(SelectExpression selectExpression)
         {
             // If the SELECT projects a single value out, we just project that with the Cosmos VALUE keyword (without VALUE,
             // Cosmos projects a JSON object containing the value).
-            if (selectExpression.UsesSingleValueProjection)
+            // TODO: Ideally, just always use VALUE for all single-projection SELECTs - but this like requires shaper changes.
+            if (selectExpression.UsesSingleValueProjection && projection is [var singleProjection])
             {
                 _sqlBuilder.Append("VALUE ");
 
-                if (projection is not [var singleProjection])
-                {
-                    throw new UnreachableException(
-                        $"Encountered SelectExpression with UsesValueProject=true and Projection.Count={projection.Count}.");
-                }
-
                 Visit(singleProjection.Expression);
             }
             // Otherwise, we'll project a JSON object; Cosmos has two syntaxes for doing so:
@@ -319,16 +314,19 @@ protected override Expression VisitSelect(SelectExpression selectExpression)
             _sqlBuilder.Append('1');
         }
 
-        if (selectExpression.Sources.Count > 0)
+        var sources = selectExpression.Sources;
+        if (sources.Count > 0)
         {
-            if (selectExpression.Sources.Count > 1)
-            {
-                throw new NotImplementedException("JOINs not yet supported");
-            }
-
             _sqlBuilder.AppendLine().Append("FROM ");
 
-            Visit(selectExpression.Sources[0]);
+            Visit(sources[0]);
+
+            for (var i = 1; i < sources.Count; i++)
+            {
+                _sqlBuilder.AppendLine().Append("JOIN ");
+
+                Visit(sources[i]);
+            }
         }
 
         if (selectExpression.Predicate != null)
@@ -752,11 +750,11 @@ protected sealed override Expression VisitSource(SourceExpression sourceExpressi
                 .Append(" IN ");
 
 
-            VisitContainerExpression(sourceExpression.ContainerExpression);
+            VisitContainerExpression(sourceExpression.Expression);
         }
         else
         {
-            VisitContainerExpression(sourceExpression.ContainerExpression);
+            VisitContainerExpression(sourceExpression.Expression);
 
             if (sourceExpression.Alias is not null)
             {
@@ -795,7 +793,7 @@ void VisitContainerExpression(Expression containerExpression)
                 }
             }
 
-            Visit(sourceExpression.ContainerExpression);
+            Visit(sourceExpression.Expression);
 
             if (subquery)
             {

src/EFCore.Cosmos/Query/Internal/CosmosQueryUtils.cs

@@ -191,7 +191,7 @@ public static bool TryExtractBareArray(
                     {
                         WithIn: true,
                         Alias: var sourceAlias,
-                        ContainerExpression: SelectExpression
+                        Expression: SelectExpression
                         {
                             Sources: [],
                             Predicate: null,
@@ -212,7 +212,7 @@ public static bool TryExtractBareArray(
 
             // For properties: SELECT i FROM i IN c.SomeArray
             // So just match any SelectExpression with IN.
-            case { Sources: [{ WithIn: true, ContainerExpression: var a, Alias: var sourceAlias }] }
+            case { Sources: [{ WithIn: true, Expression: var a, Alias: var sourceAlias }] }
                 when projectedReferenceName == sourceAlias:
             {
                 array = a;

src/EFCore.Cosmos/Query/Internal/CosmosQueryableMethodTranslatingExpressionVisitor.cs

@@ -22,7 +22,7 @@ public class CosmosQueryableMethodTranslatingExpressionVisitor : QueryableMethod
     private readonly IMethodCallTranslatorProvider _methodCallTranslatorProvider;
     private readonly CosmosSqlTranslatingExpressionVisitor _sqlTranslator;
     private readonly CosmosProjectionBindingExpressionVisitor _projectionBindingExpressionVisitor;
-    private readonly bool _subquery;
+    private bool _subquery;
     private ReadItemInfo? _readItemExpression;
 
     /// <summary>
@@ -237,7 +237,7 @@ static bool ExtractPartitionKeyFromPredicate(
     protected override Expression VisitMethodCall(MethodCallExpression methodCallExpression)
     {
         var method = methodCallExpression.Method;
-        if (method.DeclaringType == typeof(Queryable))
+        if (method.DeclaringType == typeof(Queryable) && method.IsGenericMethod)
         {
             switch (methodCallExpression.Method.Name)
             {
@@ -370,7 +370,7 @@ private ShapedQueryExpression CreateShapedQueryExpression(IEntityType entityType
             new StructuralTypeShaperExpression(
                 entityType,
                 new ProjectionBindingExpression(queryExpression, new ProjectionMember(), typeof(ValueBuffer)),
-                false));
+                nullable: false));
     }
 
     private ShapedQueryExpression CreateShapedQueryExpression(SelectExpression select, Type elementClrType)
@@ -994,7 +994,7 @@ protected override ShapedQueryExpression TranslateSelect(ShapedQueryExpression s
             return null!;
         }
 
-        var newSelectorBody = ReplacingExpressionVisitor.Replace(selector.Parameters.Single(), source.ShaperExpression, selector.Body);
+        var newSelectorBody = RemapLambdaBody(source, selector);
 
         return source.UpdateShaperExpression(_projectionBindingExpressionVisitor.Translate(selectExpression, newSelectorBody));
     }
@@ -1009,16 +1009,53 @@ protected override ShapedQueryExpression TranslateSelect(ShapedQueryExpression s
         ShapedQueryExpression source,
         LambdaExpression collectionSelector,
         LambdaExpression resultSelector)
-        => null;
+    {
+        var collectionSelectorBody = RemapLambdaBody(source, collectionSelector);
 
-    /// <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>
+        // The collection selector gets translated in subquery context; specifically, if an uncorrelated SelectMany() is attempted
+        // (from b in context.Blogs from p in context.Posts...), we want to detect that and fail translation as an uncorrelated query
+        // (see VisitExtension visitation for EntityQueryRootExpression)
+        var previousSubquery = _subquery;
+        _subquery = true;
+        try
+        {
+            if (Visit(collectionSelectorBody) is ShapedQueryExpression inner)
+            {
+                var select = (SelectExpression)source.QueryExpression;
+                var shaper = select.AddJoin(inner, source.ShaperExpression);
+
+                return TranslateTwoParameterSelector(source.UpdateShaperExpression(shaper), resultSelector);
+            }
+
+            return null;
+        }
+        finally
+        {
+            _subquery = previousSubquery;
+        }
+    }
+
+    /// <inheritdoc />
     protected override ShapedQueryExpression? TranslateSelectMany(ShapedQueryExpression source, LambdaExpression selector)
-        => null;
+    {
+        // TODO: Note that we currently never actually seem to get SelectMany without a result selector, because nav expansion rewrites
+        // that to a more complex variant with a result selector (see https://github.com/dotnet/efcore/issues/32957#issuecomment-2170950767)
+        // blogs.SelectMany(c => c.Ints) becomes:
+        // blogs
+        //     .SelectMany(p => Property(p, "Ints").AsQueryable(), (p, c) => new TransparentIdentifier`2(Outer = p, Inner = c))
+        //     .Select(ti => ti.Inner)
+
+        // TODO: In Cosmos, we currently always add a predicate for the discriminator (unless HasNoDiscriminator is explicitly specified),
+        // so the source is almost never a bare array.
+        // If we stop doing that (see #34005, #20268), and we remove the result selector problem (see just above), we should check if the
+        // source is a bare array, and simply return the ShapedQueryExpression returned from visiting the collection selector. This would
+        // remove the extra unneeded JOIN we'd currently generate.
+        var innerParameter = Expression.Parameter(selector.ReturnType.GetSequenceType(), "i");
+        var resultSelector = Expression.Lambda(
+            innerParameter, Expression.Parameter(source.Type.GetSequenceType()), innerParameter);
+
+        return TranslateSelectMany(source, selector, resultSelector);
+    }
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -1691,14 +1728,11 @@ when methodCallExpression.TryGetIndexerArguments(_queryCompilationContext.Model,
     private SqlExpression? TranslateLambdaExpression(
         ShapedQueryExpression shapedQueryExpression,
         LambdaExpression lambdaExpression)
-    {
-        var lambdaBody = RemapLambdaBody(shapedQueryExpression.ShaperExpression, lambdaExpression);
-
-        return TranslateExpression(lambdaBody);
-    }
+        => TranslateExpression(RemapLambdaBody(shapedQueryExpression, lambdaExpression));
 
-    private static Expression RemapLambdaBody(Expression shaperBody, LambdaExpression lambdaExpression)
-        => ReplacingExpressionVisitor.Replace(lambdaExpression.Parameters.Single(), shaperBody, lambdaExpression.Body);
+    private Expression RemapLambdaBody(ShapedQueryExpression shapedQueryExpression, LambdaExpression lambdaExpression)
+        => ReplacingExpressionVisitor.Replace(
+            lambdaExpression.Parameters.Single(), shapedQueryExpression.ShaperExpression, lambdaExpression.Body);
 
     private static ShapedQueryExpression AggregateResultShaper(
         ShapedQueryExpression source,
@@ -1743,4 +1777,28 @@ private static ShapedQueryExpression AggregateResultShaper(
 
         return source.UpdateShaperExpression(shaper);
     }
+
+    private ShapedQueryExpression TranslateTwoParameterSelector(ShapedQueryExpression source, LambdaExpression resultSelector)
+    {
+        var transparentIdentifierType = source.ShaperExpression.Type;
+        var transparentIdentifierParameter = Expression.Parameter(transparentIdentifierType);
+
+        Expression original1 = resultSelector.Parameters[0];
+        var replacement1 = AccessField(transparentIdentifierType, transparentIdentifierParameter, "Outer");
+        Expression original2 = resultSelector.Parameters[1];
+        var replacement2 = AccessField(transparentIdentifierType, transparentIdentifierParameter, "Inner");
+        var newResultSelector = Expression.Lambda(
+            new ReplacingExpressionVisitor(
+                    new[] { original1, original2 }, new[] { replacement1, replacement2 })
+                .Visit(resultSelector.Body),
+            transparentIdentifierParameter);
+
+        return TranslateSelect(source, newResultSelector);
+    }
+
+    private static Expression AccessField(
+        Type transparentIdentifierType,
+        Expression targetExpression,
+        string fieldName)
+        => Expression.Field(targetExpression, transparentIdentifierType.GetTypeInfo().GetDeclaredField(fieldName)!);
 }