Skip to content

Commit dd1f1bf

Browse files
BorisDogBorisDog
authored
CSHARP-4851: Render full form for $and/$eq in $vectorSearch stage filter (#1220)
1 parent 3274297 commit dd1f1bf

File tree

6 files changed

+197
-7
lines changed

6 files changed

+197
-7
lines changed

‎src/MongoDB.Driver/FilterDefinition.cs

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
using System;
1717
using System.Linq.Expressions;
18+
using System.Threading;
1819
using MongoDB.Bson;
1920
using MongoDB.Bson.Serialization;
2021
using MongoDB.Driver.Core.Misc;
@@ -285,4 +286,30 @@ public override BsonDocument Render(IBsonSerializer<TDocument> documentSerialize
285286
return new BsonDocumentWrapper(_obj, serializer);
286287
}
287288
}
289+
290+
internal static class FilterDefinitionRenderContext
291+
{
292+
private static readonly AsyncLocal<bool> __renderDollarForm = new AsyncLocal<bool>();
293+
294+
public static bool RenderDollarForm
295+
{
296+
get => __renderDollarForm.Value;
297+
set => __renderDollarForm.Value = value;
298+
}
299+
300+
public static IDisposable StartRender(bool renderDollarForm) => new FilterDefinitionRenderContextDisposer(renderDollarForm);
301+
302+
private sealed class FilterDefinitionRenderContextDisposer : IDisposable
303+
{
304+
public FilterDefinitionRenderContextDisposer(bool renderDollarForm)
305+
{
306+
RenderDollarForm = renderDollarForm;
307+
}
308+
309+
public void Dispose()
310+
{
311+
RenderDollarForm = false;
312+
}
313+
}
314+
}
288315
}

‎src/MongoDB.Driver/FilterDefinitionBuilder.cs

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1706,6 +1706,15 @@ public override BsonDocument Render(IBsonSerializer<TDocument> documentSerialize
17061706
}
17071707
}
17081708

1709+
if (FilterDefinitionRenderContext.RenderDollarForm)
1710+
{
1711+
if (document.ElementCount > 1 ||
1712+
document.ElementCount == 1 && document.GetElement(0).Name != "$and")
1713+
{
1714+
PromoteFilterToDollarForm(document, null);
1715+
}
1716+
}
1717+
17091718
return document;
17101719
}
17111720

@@ -1759,14 +1768,19 @@ existingClause.Value is BsonDocument existingClauseValue &&
17591768
static bool IsFieldName(string fieldOrOperatorName) => !fieldOrOperatorName.StartsWith("$");
17601769
}
17611770

1762-
private static void PromoteFilterToDollarForm(BsonDocument document, BsonElement clause)
1771+
private static void PromoteFilterToDollarForm(BsonDocument document, BsonElement? clause)
17631772
{
17641773
var clauses = new BsonArray();
17651774
foreach (var queryElement in document)
17661775
{
17671776
clauses.Add(new BsonDocument(queryElement));
17681777
}
1769-
clauses.Add(new BsonDocument(clause));
1778+
1779+
if (clause != null)
1780+
{
1781+
clauses.Add(new BsonDocument(clause.Value));
1782+
}
1783+
17701784
document.Clear();
17711785
document.Add("$and", clauses);
17721786
}
@@ -2339,7 +2353,19 @@ public override BsonDocument Render(IBsonSerializer<TDocument> documentSerialize
23392353
var context = BsonSerializationContext.CreateRoot(bsonWriter);
23402354
bsonWriter.WriteStartDocument();
23412355
bsonWriter.WriteName(renderedField.FieldName);
2342-
renderedField.ValueSerializer.Serialize(context, _value);
2356+
2357+
if (FilterDefinitionRenderContext.RenderDollarForm)
2358+
{
2359+
bsonWriter.WriteStartDocument();
2360+
bsonWriter.WriteName("$eq");
2361+
renderedField.ValueSerializer.Serialize(context, _value);
2362+
bsonWriter.WriteEndDocument();
2363+
}
2364+
else
2365+
{
2366+
renderedField.ValueSerializer.Serialize(context, _value);
2367+
}
2368+
23432369
bsonWriter.WriteEndDocument();
23442370
}
23452371

‎src/MongoDB.Driver/PipelineStageDefinitionBuilder.cs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1981,14 +1981,20 @@ public static PipelineStageDefinition<TInput, TInput> VectorSearch<TInput>(
19811981
{ "limit", limit },
19821982
{ "numCandidates", options?.NumberOfCandidates ?? limit * 10 },
19831983
{ "index", options?.IndexName ?? "default" },
1984-
{ "filter", () => options?.Filter?.Render(s, sr, linqProvider), options?.Filter != null },
1984+
{ "filter", () => RenderFilter(s, sr, linqProvider), options?.Filter != null },
19851985
};
19861986

19871987
var document = new BsonDocument(operatorName, vectorSearchOperator);
19881988
return new RenderedPipelineStageDefinition<TInput>(operatorName, document, s);
19891989
});
19901990

19911991
return stage;
1992+
1993+
BsonDocument RenderFilter(IBsonSerializer<TInput> documentSerializer, IBsonSerializerRegistry serializerRegistry, LinqProvider linqProvider)
1994+
{
1995+
using var renderContext = FilterDefinitionRenderContext.StartRender(true);
1996+
return options.Filter.Render(documentSerializer, serializerRegistry, linqProvider);
1997+
}
19921998
}
19931999

19942000
// private methods

‎tests/MongoDB.Driver.Tests/FilterDefinitionRenderContextTests.cs

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
/* Copyright 2010-present MongoDB Inc.
2+
*
3+
* Licensed under the Apache License, Version 2.0 (the "License");
4+
* you may not use this file except in compliance with the License.
5+
* You may obtain a copy of the License at
6+
*
7+
* http://www.apache.org/licenses/LICENSE-2.0
8+
*
9+
* Unless required by applicable law or agreed to in writing, software
10+
* distributed under the License is distributed on an "AS IS" BASIS,
11+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
* See the License for the specific language governing permissions and
13+
* limitations under the License.
14+
*/
15+
16+
using System.Collections.Generic;
17+
using System.Threading;
18+
using System.Threading.Tasks;
19+
using FluentAssertions;
20+
using MongoDB.Bson;
21+
using MongoDB.Bson.Serialization;
22+
using Xunit;
23+
24+
namespace MongoDB.Driver.Tests
25+
{
26+
public class FilterDefinitionRenderContextTests
27+
{
28+
[Fact]
29+
public void FilterDefinitionRenderContext_dollarForm_default_value_should_be_false()
30+
{
31+
FilterDefinitionRenderContext.RenderDollarForm.Should().BeFalse();
32+
}
33+
34+
[Fact]
35+
public async Task FilterDefinitionRenderContext_should_be_scoped_to_task()
36+
{
37+
bool? renderDollarFormValueObservedByTask = null;
38+
var taskReadyToRenderEvent = new ManualResetEventSlim();
39+
var unblockTaskEvent = new ManualResetEventSlim();
40+
var rendererTask = Task.Run(RendererTask);
41+
42+
// Wait for task to set RenderDollarForm = true;
43+
taskReadyToRenderEvent.Wait();
44+
45+
// Try to 'override' RenderDollarForm value and unblock the task
46+
FilterDefinitionRenderContext.RenderDollarForm = false;
47+
unblockTaskEvent.Set();
48+
49+
// Wait fot task to finish and set renderDollarFormValueObservedByTask
50+
await rendererTask;
51+
52+
renderDollarFormValueObservedByTask.Should().Be(true);
53+
FilterDefinitionRenderContext.RenderDollarForm.Should().Be(false);
54+
55+
void RendererTask()
56+
{
57+
using var renderContext = FilterDefinitionRenderContext.StartRender(true);
58+
59+
taskReadyToRenderEvent.Set();
60+
unblockTaskEvent.Wait();
61+
62+
renderDollarFormValueObservedByTask = FilterDefinitionRenderContext.RenderDollarForm;
63+
}
64+
}
65+
66+
[Theory]
67+
[MemberData(nameof(Correct_form_should_be_rendered_test_cases))]
68+
public void Correct_form_should_be_rendered(FilterDefinition<BsonDocument> filterDefinition, bool renderDollarForm, string expectedFilter)
69+
{
70+
var serializerRegistry = BsonSerializer.SerializerRegistry;
71+
var documentSerializer = serializerRegistry.GetSerializer<BsonDocument>();
72+
var expectedFilterDocument = BsonDocument.Parse(expectedFilter);
73+
74+
using var renderContext = FilterDefinitionRenderContext.StartRender(renderDollarForm);
75+
var actualFilter = filterDefinition.Render(documentSerializer, serializerRegistry);
76+
77+
actualFilter.Should().Be(expectedFilterDocument);
78+
}
79+
80+
public static IEnumerable<object[]> Correct_form_should_be_rendered_test_cases()
81+
{
82+
return new object[][]
83+
{
84+
// $eq
85+
new object[] { Eq("a", 1), false, "{ a : 1 }" },
86+
new object[] { Eq("a", 1), true, "{ a: { $eq: 1 } }" },
87+
88+
// $and
89+
new object[] { Gt("a", 1) & Gt("b", 2), false, "{ a : { $gt : 1 }, b : { $gt : 2 } }" },
90+
new object[] { Gt("a", 1) & Gt("b", 2), true, "{ $and: [ { a : { $gt : 1 }}, { b : { $gt : 2 }} ] }" },
91+
new object[] { Gt("a", 1) & Gt("b", 2) & Gt("c", 3), false, "{ a : { $gt : 1 }, b : { $gt : 2 }, c : { $gt : 3 } }" },
92+
new object[] { Gt("a", 1) & Gt("b", 2) & Gt("c", 3), true, "{ $and: [ { a : { $gt : 1 }}, { b : { $gt : 2 }}, { c : { $gt : 3 }} ] }" },
93+
94+
// nested $eq
95+
new object[] { Eq("a", 1) | Eq("b", 2), false, "{ $or: [{ a : 1 }, { b : 2 }] }" },
96+
new object[] { Eq("a", 1) | Eq("b", 2), true, "{ $or: [{ a : { $eq: 1 }}, { b : { $eq : 2 }}] }" },
97+
new object[] { !(Eq("a", 1) | Eq("b", 2)), false, "{ $nor: [ { a : 1 }, { b : 2 }] }" },
98+
new object[] { !(Eq("a", 1) | Eq("b", 2)), true, "{ $nor: [{ a : { $eq: 1 }}, { b : { $eq: 2 }}] }" },
99+
100+
// nested $and
101+
new object[] { Gt("a", 1) & Gt("b", 2) | Gt("c", 3) & Gt("d", 4), false, "{ $or: [{ a : { $gt : 1 }, b : { $gt : 2 }}, { c : { $gt : 3 }, d : { $gt : 4 }}] }" },
102+
new object[] { Gt("a", 1) & Gt("b", 2) | Gt("c", 3) & Gt("d", 4), true, "{ $or: [{ $and: [ { a : { $gt : 1 }}, { b : { $gt : 2 }} ] }, { $and: [ { c : { $gt : 3 }}, { d : { $gt : 4 }} ] }] }" },
103+
104+
// $eq and $and
105+
new object[] { Eq("a", 1) & Eq("b", 2), false, "{ a : 1, b : 2 }" },
106+
new object[] { Eq("a", 1) & Eq("b", 2), true, "{ $and: [{ a: { $eq: 1 }}, { b: { $eq: 2 }}] }" },
107+
108+
// $eq and $and nested
109+
new object[] { !(Eq("a", 1) & Eq("b", 2)), false, "{ $nor: [{ a : 1, b : 2 }] }" },
110+
new object[] { !(Eq("a", 1) & Eq("b", 2)), true, "{ $nor: [{ $and: [{ a: { $eq: 1 }}, { b: { $eq: 2 }}] }] }" },
111+
112+
// always dollar form
113+
new object[] { Eq("a", 1) & Eq("a", 2), false, "{ $and : [{ a : 1 }, { a : 2 }] }" },
114+
new object[] { Eq("a", 1) & Eq("a", 2), true, "{ $and : [{ a : { $eq: 1 } }, { a : { $eq: 2 } }] }" },
115+
new object[] { Gt("a", 1) & Gt("a", 2), false, "{ $and : [{ a : { $gt : 1 } }, { a : { $gt : 2 } }] }" },
116+
new object[] { Gt("a", 1) & Gt("a", 2), true, "{ $and : [{ a : { $gt : 1 } }, { a : { $gt : 2 } }] }" },
117+
};
118+
}
119+
120+
private static FilterDefinition<BsonDocument> Eq(string field, int value) => GetBuilder().Eq(field, value);
121+
private static FilterDefinition<BsonDocument> Gt(string field, int value) => GetBuilder().Gt(field, value);
122+
123+
private static FilterDefinitionBuilder<BsonDocument> GetBuilder() => new();
124+
}
125+
}

‎tests/MongoDB.Driver.Tests/PipelineDefinitionBuilderTests.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -416,14 +416,14 @@ public void VectorSearch_should_add_expected_stage_with_options()
416416
var pipeline = new EmptyPipelineDefinition<BsonDocument>();
417417
var options = new VectorSearchOptions<BsonDocument>()
418418
{
419-
Filter = Builders<BsonDocument>.Filter.Eq("y", "val"),
419+
Filter = Builders<BsonDocument>.Filter.Eq("x", 1) & Builders<BsonDocument>.Filter.Eq("y", 2),
420420
IndexName = "index_name",
421421
NumberOfCandidates = 123
422422
};
423423
var result = pipeline.VectorSearch("x", new[] { 1.0, 2.0, 3.0 }, 1, options);
424424

425425
var stages = RenderStages(result, BsonDocumentSerializer.Instance);
426-
stages[0].Should().Be("{ $vectorSearch: { queryVector: [1.0, 2.0, 3.0], path: 'x', limit: 1, numCandidates: 123, index: 'index_name', filter: { y: 'val' } } }");
426+
stages[0].Should().Be("{ $vectorSearch: { queryVector: [1.0, 2.0, 3.0], path: 'x', limit: 1, numCandidates: 123, index: 'index_name', filter : { $and : [{ x : { $eq : 1 } }, { y : { $eq : 2 } }] } } }");
427427
}
428428

429429
[Fact]

‎tests/MongoDB.Driver.Tests/Search/AtlasSearchTests.cs

Lines changed: 7 additions & 1 deletion
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)