Skip to content

Commit f80e5f9

Browse files
authored
.Net: Fix batch operation non-string keys with MongoDB (#13036)
### Motivation and Context The ObjectID and Guid string support work in #12827 did not include support for the method overloads that take an IEnumerable of keys. I suspect this was due to merge/out of sync branches and was then missed with the tests being disabled (they now pass). ### Description Adds the support for Guid and ObjectID keys to the GetAsync, DeleteAsync and UpdateAsync methods that take an IEnumerable. ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄
1 parent 0ec81c8 commit f80e5f9

File tree

1 file changed

+10
-11
lines changed

1 file changed

+10
-11
lines changed

dotnet/src/VectorData/MongoDB/MongoCollection.cs

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -171,9 +171,7 @@ public override async Task DeleteAsync(IEnumerable<TKey> keys, CancellationToken
171171
{
172172
Verify.NotNull(keys);
173173

174-
var stringKeys = keys is IEnumerable<string> k ? k : keys.Cast<string>();
175-
176-
await this.RunOperationAsync("DeleteMany", () => this._mongoCollection.DeleteManyAsync(this.GetFilterByIds(stringKeys), cancellationToken))
174+
await this.RunOperationAsync("DeleteMany", () => this._mongoCollection.DeleteManyAsync(this.GetFilterByIds(keys), cancellationToken))
177175
.ConfigureAwait(false);
178176
}
179177

@@ -220,10 +218,8 @@ public override async IAsyncEnumerable<TRecord> GetAsync(
220218
throw new NotSupportedException(VectorDataStrings.IncludeVectorsNotSupportedWithEmbeddingGeneration);
221219
}
222220

223-
var stringKeys = keys is IEnumerable<string> k ? k : keys.Cast<string>();
224-
225221
using var cursor = await this
226-
.FindAsync(this.GetFilterByIds(stringKeys), top: null, skip: null, includeVectors, sortDefinition: null, cancellationToken)
222+
.FindAsync(this.GetFilterByIds(keys), top: null, skip: null, includeVectors, sortDefinition: null, cancellationToken)
227223
.ConfigureAwait(false);
228224

229225
while (await cursor.MoveNextAsync(cancellationToken).ConfigureAwait(false))
@@ -270,14 +266,17 @@ private async Task UpsertCoreAsync(TRecord record, int recordIndex, IReadOnlyLis
270266
var replaceOptions = new ReplaceOptions { IsUpsert = true };
271267
var storageModel = this._mapper.MapFromDataToStorageModel(record, recordIndex, generatedEmbeddings);
272268

273-
var key = storageModel[MongoConstants.MongoReservedKeyPropertyName];
269+
var key = GetStorageKey(storageModel);
274270

275271
await this.RunOperationAsync(OperationName, async () =>
276272
await this._mongoCollection
277273
.ReplaceOneAsync(this.GetFilterById(key), storageModel, replaceOptions, cancellationToken)
278274
.ConfigureAwait(false)).ConfigureAwait(false);
279275
}
280276

277+
private static TKey GetStorageKey(BsonDocument document)
278+
=> (TKey)BsonTypeMapper.MapToDotNetValue(document[MongoConstants.MongoReservedKeyPropertyName]);
279+
281280
private static async ValueTask<(IEnumerable<TRecord> records, IReadOnlyList<Embedding>?[]?)> ProcessEmbeddingsAsync(
282281
CollectionModel model,
283282
IEnumerable<TRecord> records,
@@ -676,11 +675,11 @@ private async IAsyncEnumerable<VectorSearchResult<TRecord>> EnumerateAndMapSearc
676675
}
677676
}
678677

679-
private FilterDefinition<BsonDocument> GetFilterById(object id)
680-
=> Builders<BsonDocument>.Filter.Eq(document => document[MongoConstants.MongoReservedKeyPropertyName], id);
678+
private FilterDefinition<BsonDocument> GetFilterById(TKey id)
679+
=> Builders<BsonDocument>.Filter.Eq(MongoConstants.MongoReservedKeyPropertyName, id);
681680

682-
private FilterDefinition<BsonDocument> GetFilterByIds(IEnumerable<object> ids)
683-
=> Builders<BsonDocument>.Filter.In(document => document[MongoConstants.MongoReservedKeyPropertyName], ids);
681+
private FilterDefinition<BsonDocument> GetFilterByIds(IEnumerable<TKey> ids)
682+
=> Builders<BsonDocument>.Filter.In(MongoConstants.MongoReservedKeyPropertyName, ids);
684683

685684
private async Task<bool> InternalCollectionExistsAsync(CancellationToken cancellationToken)
686685
{

0 commit comments

Comments
 (0)