Skip to content

Support conflicting partition key predicates in Cosmos queries #38240

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lorenacr wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Support conflicting partition key predicates in Cosmos queries #38240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ee3c440
Support conflicting partition key predicates in Cosmos queries
lorenacr May 7, 2026
5ad8aa4
Enhance partition key conflict validation in Cosmos queries
lorenacr May 8, 2026
dca176f
Add using directive for internal Cosmos functionality in ReadItemPart…
lorenacr May 8, 2026
056d94d
Enhance partition key conflict handling in Cosmos queries
lorenacr May 12, 2026
805c185
Refactor ReadItem optimization logic in Cosmos queries
lorenacr May 14, 2026
81f8ba6
Refactor handling of conflicting partition key predicates in Cosmos q…
lorenacr May 27, 2026
f073862
Merge upstream/main into Issue38238
lorenacr May 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs
View file
Open in desktop

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

2 changes: 1 addition & 1 deletion src/EFCore.Cosmos/Properties/CosmosStrings.resx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
<?xml version="1.0" encoding="utf-8"?>
<?xml version="1.0" encoding="utf-8"?>
<root>
<!--
Microsoft ResX Schema
Expand Down
60 changes: 38 additions & 22 deletions src/EFCore.Cosmos/Query/Internal/CosmosReadItemAndPartitionKeysExtractor.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,22 @@ public virtual Expression ExtractPartitionKeysAndId(

_rootAlias = rootSource.Alias;

Expression UnwrapShaperForReadItem(Expression shaper)
Copy link
Copy Markdown
Member

@roji roji May 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason for this rename and move of Unwrap()? At least keep it in the same place to reduce the diff?

{
if (shaper is UnaryExpression { NodeType: ExpressionType.Convert } convert
&& convert.Type == typeof(object))
{
shaper = convert.Operand;
}

while (shaper is IncludeExpression { EntityExpression: var nested })
{
shaper = nested;
}

return shaper;
}

// We're going to be looking for equality comparisons on the JSON id definition properties and the partition key properties of the
// entity type; build a dictionary where the properties are the keys, and where the values are expressions that will get populated
// from the tree (either constants or parameters).
Expand Down Expand Up @@ -88,6 +104,19 @@ public virtual Expression ExtractPartitionKeysAndId(
var allIdPropertiesSpecified =
_jsonIdPropertyValues.Values.All(p => p is not null) && _jsonIdPropertyValues.Count > 0;

// WithPartitionKey will clear _partitionKeyPropertyValues during the lift pass below; snapshot predicate partition key
// comparisons first so we can avoid ReadItem when both WithPartitionKey and the predicate specify partition keys (see #38238).
var hadWithPartitionKey = queryCompilationContext.PartitionKeyPropertyValues.Count > 0;
Dictionary<IProperty, (Expression? ValueExpression, Expression? OriginalExpression)>? predicatePartitionKeySnapshot = null;
if (hadWithPartitionKey)
{
predicatePartitionKeySnapshot = new Dictionary<IProperty, (Expression?, Expression?)>(_partitionKeyPropertyValues);
}

var predicateSpecifiesPartitionKey = hadWithPartitionKey
&& predicatePartitionKeySnapshot is not null
&& partitionKeyProperties.Any(p => predicatePartitionKeySnapshot[p].ValueExpression is not null);

// First, go over the partition key properties and lift them from the predicate to the query compilation context, as possible.
// We do this only as long as all partition key values are provided; the moment there's a gap we stop (so if PK1 and PK3 are
// provided but not PK2, only PK1 will be lifted out).
Expand All @@ -111,20 +140,23 @@ public virtual Expression ExtractPartitionKeysAndId(
}
}

// Now, attempt to also transform the query to ReadItem form; this is only possible if all JSON ID properties were compared in the
// predicate, and *all* partition key values are specified(in the predicate or via WithPartitionKey)
if (_isPredicateCompatibleWithReadItem
var willUseReadItemOptimization = _isPredicateCompatibleWithReadItem
&& allIdPropertiesSpecified
&& queryCompilationContext.PartitionKeyPropertyValues.Count == partitionKeyProperties.Count
&& !predicateSpecifiesPartitionKey
&& select is
{
Offset: null or SqlConstantExpression { Value: 0 },
Limit: null or SqlConstantExpression { Value: > 0 }
}
// We only transform to ReadItem if the entire document (i.e. root entity type) is being projected out.
// Using ReadItem even when a projection is present is tracked by #34163.
&& Unwrap(shapedQuery.ShaperExpression) is StructuralTypeShaperExpression { StructuralType: var projectedStructuralType }
&& projectedStructuralType == _entityType)
&& UnwrapShaperForReadItem(shapedQuery.ShaperExpression) is StructuralTypeShaperExpression { StructuralType: var projectedStructuralType }
&& projectedStructuralType == _entityType;

// Now, attempt to also transform the query to ReadItem form; this is only possible if all JSON ID properties were compared in the
// predicate, and *all* partition key values are specified(in the predicate or via WithPartitionKey)
if (willUseReadItemOptimization)
{
return shapedQuery.UpdateQueryExpression(select.WithReadItemInfo(new ReadItemInfo(_jsonIdPropertyValues!)));
}
Expand Down Expand Up @@ -152,22 +184,6 @@ public virtual Expression ExtractPartitionKeysAndId(
}

return shapedQuery;

Expression Unwrap(Expression shaper)
{
if (shaper is UnaryExpression { NodeType: ExpressionType.Convert } convert
&& convert.Type == typeof(object))
{
shaper = convert.Operand;
}

while (shaper is IncludeExpression { EntityExpression: var nested })
{
shaper = nested;
}

return shaper;
}
}

/// <summary>
Expand Down Expand Up @@ -320,7 +336,7 @@ void ProcessPropertyComparison(string propertyName, SqlExpression propertyValue,
// call. Note that this is always considered a compatible comparison for ReadItem.
if (propertyName == property.GetJsonPropertyName()
&& _partitionKeyPropertyValues.TryGetValue(property, out var previousValues)
&& (previousValues.ValueExpression is null || previousValues.Equals(propertyValue)))
&& (previousValues.ValueExpression is null || previousValues.ValueExpression.Equals(propertyValue)))
{
_partitionKeyPropertyValues[property] = (ValueExpression: propertyValue, OriginalExpression: originalExpression);
return;
Expand Down
55 changes: 53 additions & 2 deletions test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryDiscriminatorInIdTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -445,6 +445,47 @@ FROM root c
""");
}

public override async Task WithPartitionKey_and_conflicting_partition_key_predicate_returns_empty()
{
await base.WithPartitionKey_and_conflicting_partition_key_predicate_returns_empty();

AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = "PK2")))
""");
}

public override async Task WithPartitionKey_and_partition_key_predicate_same_local()
{
await base.WithPartitionKey_and_partition_key_predicate_same_local();

AssertSql(
"""
@partitionKey='PK1'

SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @partitionKey)))
""");
}

public override async Task WithPartitionKey_and_partition_key_predicate_two_locals_same_value()
{
await base.WithPartitionKey_and_partition_key_predicate_two_locals_same_value();

AssertSql(
"""
@pkForWith='PK1'
@pkForWhere='PK1'

SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @pkForWhere)))
""");
}

public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
{
await base.ReadItem_with_WithPartitionKey_with_only_partition_key();
Expand Down Expand Up @@ -797,7 +838,12 @@ public override async Task Both_WithPartitionKey_and_predicate_comparisons_with_
{
await base.Both_WithPartitionKey_and_predicate_comparisons_with_same_values_with_only_partition_key_leaf();

AssertSql("""ReadItem(["PK1c"], DerivedOnlySinglePartitionKeyEntity|PK1c)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["$type"] = "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1c"))
""");
}

public override async Task ReadItem_with_hierarchical_partition_key_leaf()
Expand Down Expand Up @@ -891,7 +937,12 @@ public override async Task ReadItem_with_WithPartitionKey_with_only_partition_ke
{
await base.ReadItem_with_WithPartitionKey_with_only_partition_key_leaf();

AssertSql("""ReadItem(["PK1c"], DerivedOnlySinglePartitionKeyEntity|PK1c)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["$type"] = "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1c"))
""");
}

public override async Task Multiple_incompatible_predicate_comparisons_cause_no_ReadItem_leaf()
Expand Down
69 changes: 65 additions & 4 deletions test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryNoDiscriminatorInIdTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,7 +275,12 @@ public override async Task Both_WithPartitionKey_and_predicate_comparisons_with_
{
await base.Both_WithPartitionKey_and_predicate_comparisons_with_same_values_with_only_partition_key();

AssertSql("""ReadItem(["PK1a"], PK1a)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE (c["PartitionKey"] = "PK1a")
""");
}

public override async Task ReadItem_with_hierarchical_partition_key()
Expand Down Expand Up @@ -362,11 +367,57 @@ public override async Task ReadItem_with_WithPartitionKey()
AssertSql("""ReadItem(["PK1"], b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
}

public override async Task WithPartitionKey_and_conflicting_partition_key_predicate_returns_empty()
{
await base.WithPartitionKey_and_conflicting_partition_key_predicate_returns_empty();

AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = "PK2"))
""");
}

public override async Task WithPartitionKey_and_partition_key_predicate_same_local()
{
await base.WithPartitionKey_and_partition_key_predicate_same_local();

AssertSql(
"""
@partitionKey='PK1'

SELECT VALUE c
FROM root c
WHERE ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @partitionKey))
""");
}

public override async Task WithPartitionKey_and_partition_key_predicate_two_locals_same_value()
{
await base.WithPartitionKey_and_partition_key_predicate_two_locals_same_value();

AssertSql(
"""
@pkForWith='PK1'
@pkForWhere='PK1'

SELECT VALUE c
FROM root c
WHERE ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @pkForWhere))
""");
}

public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
{
await base.ReadItem_with_WithPartitionKey_with_only_partition_key();

AssertSql("""ReadItem(["PK1a"], PK1a)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE (c["PartitionKey"] = "PK1a")
""");
}

public override async Task Multiple_incompatible_predicate_comparisons_cause_no_ReadItem()
Expand Down Expand Up @@ -673,7 +724,12 @@ public override async Task Both_WithPartitionKey_and_predicate_comparisons_with_
{
await base.Both_WithPartitionKey_and_predicate_comparisons_with_same_values_with_only_partition_key_leaf();

AssertSql("""ReadItem(["PK1c"], PK1c)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["$type"] = "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1c"))
""");
}

public override async Task ReadItem_with_hierarchical_partition_key_leaf()
Expand Down Expand Up @@ -764,7 +820,12 @@ public override async Task ReadItem_with_WithPartitionKey_with_only_partition_ke
{
await base.ReadItem_with_WithPartitionKey_with_only_partition_key_leaf();

AssertSql("""ReadItem(["PK1c"], PK1c)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["$type"] = "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1c"))
""");
}

public override async Task Multiple_incompatible_predicate_comparisons_cause_no_ReadItem_leaf()
Expand Down
69 changes: 65 additions & 4 deletions ...EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryRootDiscriminatorInIdTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,12 @@ public override async Task Both_WithPartitionKey_and_predicate_comparisons_with_
{
await base.Both_WithPartitionKey_and_predicate_comparisons_with_same_values_with_only_partition_key();

AssertSql("""ReadItem(["PK1a"], OnlySinglePartitionKeyEntity|PK1a)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("OnlySinglePartitionKeyEntity", "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1a"))
""");
}

public override async Task ReadItem_with_hierarchical_partition_key()
Expand Down Expand Up @@ -355,11 +360,57 @@ public override async Task ReadItem_with_WithPartitionKey()
AssertSql("""ReadItem(["PK1"], SinglePartitionKeyEntity|b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
}

public override async Task WithPartitionKey_and_conflicting_partition_key_predicate_returns_empty()
{
await base.WithPartitionKey_and_conflicting_partition_key_predicate_returns_empty();

AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = "PK2")))
""");
}

public override async Task WithPartitionKey_and_partition_key_predicate_same_local()
{
await base.WithPartitionKey_and_partition_key_predicate_same_local();

AssertSql(
"""
@partitionKey='PK1'

SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @partitionKey)))
""");
}

public override async Task WithPartitionKey_and_partition_key_predicate_two_locals_same_value()
{
await base.WithPartitionKey_and_partition_key_predicate_two_locals_same_value();

AssertSql(
"""
@pkForWith='PK1'
@pkForWhere='PK1'

SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @pkForWhere)))
""");
}

public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
{
await base.ReadItem_with_WithPartitionKey_with_only_partition_key();

AssertSql("""ReadItem(["PK1a"], OnlySinglePartitionKeyEntity|PK1a)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE (c["$type"] IN ("OnlySinglePartitionKeyEntity", "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1a"))
""");
}

public override async Task Multiple_incompatible_predicate_comparisons_cause_no_ReadItem()
Expand Down Expand Up @@ -668,7 +719,12 @@ public override async Task Both_WithPartitionKey_and_predicate_comparisons_with_
{
await base.Both_WithPartitionKey_and_predicate_comparisons_with_same_values_with_only_partition_key_leaf();

AssertSql("""ReadItem(["PK1c"], OnlySinglePartitionKeyEntity|PK1c)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["$type"] = "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1c"))
""");
}

public override async Task ReadItem_with_hierarchical_partition_key_leaf()
Expand Down Expand Up @@ -759,7 +815,12 @@ public override async Task ReadItem_with_WithPartitionKey_with_only_partition_ke
{
await base.ReadItem_with_WithPartitionKey_with_only_partition_key_leaf();

AssertSql("""ReadItem(["PK1c"], OnlySinglePartitionKeyEntity|PK1c)""");
AssertSql(
"""
SELECT VALUE c
FROM root c
WHERE ((c["$type"] = "DerivedOnlySinglePartitionKeyEntity") AND (c["PartitionKey"] = "PK1c"))
""");
}

public override async Task Multiple_incompatible_predicate_comparisons_cause_no_ReadItem_leaf()
Expand Down
Loading