Skip to content

Commit

Permalink
Cosmos: Implement SelectMany
Browse files Browse the repository at this point in the history
Also introduces substantial infrastructure for general joins

Closes #17312
  • Loading branch information
roji committed Jun 18, 2024
1 parent 0d75b75 commit 8bb7df1
Show file tree
Hide file tree
Showing 19 changed files with 666 additions and 233 deletions.
6 changes: 6 additions & 0 deletions src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions src/EFCore.Cosmos/Properties/CosmosStrings.resx
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
32 changes: 15 additions & 17 deletions src/EFCore.Cosmos/Query/Internal/CosmosQuerySqlGenerator.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)
Expand Down Expand Up @@ -728,11 +726,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)
{
Expand Down Expand Up @@ -771,7 +769,7 @@ void VisitContainerExpression(Expression containerExpression)
}
}

Visit(sourceExpression.ContainerExpression);
Visit(sourceExpression.Expression);

if (subquery)
{
Expand Down
4 changes: 2 additions & 2 deletions src/EFCore.Cosmos/Query/Internal/CosmosQueryUtils.cs
Original file line number Diff line number Diff line change
Expand Up @@ -191,7 +191,7 @@ public static bool TryExtractBareArray(
{
WithIn: true,
Alias: var sourceAlias,
ContainerExpression: SelectExpression
Expression: SelectExpression
{
Sources: [],
Predicate: null,
Expand All @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,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>
Expand Down Expand Up @@ -236,7 +236,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)
{
Expand Down Expand Up @@ -369,7 +369,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)
Expand Down Expand Up @@ -985,7 +985,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));
}
Expand All @@ -1000,16 +1000,52 @@ 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, so the source is never a bare array.
// If we stop doing that (see #34005, #20268), 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
Expand Down Expand Up @@ -1684,14 +1720,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,
Expand Down Expand Up @@ -1736,4 +1769,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)!);
}
Loading

0 comments on commit 8bb7df1

Please sign in to comment.