From 2f7c242d68d096f6f69ba48cfe4bfd673ece24a6 Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Wed, 10 Dec 2025 19:38:01 +0100 Subject: [PATCH] [MEVD] Implement key auto-generation Closes #11485 --- .../Memory/MongoDB/MongoModelBuilder.cs | 13 +- .../AzureAISearch/AzureAISearchCollection.cs | 14 +- .../AzureAISearchDynamicModelBuilder.cs | 12 +- .../AzureAISearchModelBuilder.cs | 19 +- .../CosmosMongoDB/CosmosMongoCollection.cs | 25 ++ .../CosmosNoSql/CosmosNoSqlCollection.cs | 7 + .../CosmosNoSql/CosmosNoSqlModelBuilder.cs | 10 +- .../VectorData/InMemory/InMemoryCollection.cs | 10 +- .../InMemory/InMemoryModelBuilder.cs | 12 +- .../src/VectorData/MongoDB/MongoCollection.cs | 26 ++ .../VectorData/PgVector/PostgresCollection.cs | 118 ++++----- .../PgVector/PostgresModelBuilder.cs | 21 +- .../VectorData/PgVector/PostgresSqlBuilder.cs | 174 ++++++++---- .../VectorData/Pinecone/PineconeCollection.cs | 21 +- .../Pinecone/PineconeModelBuilder.cs | 10 +- .../src/VectorData/Qdrant/QdrantCollection.cs | 15 +- .../VectorData/Qdrant/QdrantModelBuilder.cs | 10 +- .../Redis/RedisHashSetCollection.cs | 7 + .../VectorData/Redis/RedisJsonCollection.cs | 14 + .../Redis/RedisJsonDynamicModelBuilder.cs | 10 +- .../VectorData/Redis/RedisJsonModelBuilder.cs | 10 +- .../src/VectorData/Redis/RedisModelBuilder.cs | 10 +- .../SqlServer/SqlServerCollection.cs | 74 +++++- .../SqlServer/SqlServerCommandBuilder.cs | 213 ++++++++------- .../SqlServer/SqlServerModelBuilder.cs | 16 +- .../VectorData/SqliteVec/SqliteCollection.cs | 214 +++++++-------- .../SqliteVec/SqliteCommandBuilder.cs | 107 +++++--- .../SqliteVec/SqliteModelBuilder.cs | 12 +- .../CollectionModelBuilder.cs | 24 +- .../ProviderServices/KeyPropertyModel.cs | 7 +- .../VectorStoreKeyAttribute.cs | 8 + .../VectorStoreKeyProperty.cs | 8 + .../VectorData/Weaviate/WeaviateCollection.cs | 13 +- .../Weaviate/WeaviateModelBuilder.cs | 10 +- .../TypeTests/CosmosMongoKeyTypeTests.cs | 6 +- .../TypeTests/InMemoryKeyTypeTests.cs | 6 +- .../TypeTests/MongoKeyTypeTests.cs | 6 +- .../Support/PostgresTestStore.cs | 2 +- .../TypeTests/PostgresKeyTypeTests.cs | 4 +- .../PostgresSqlBuilderTests.cs | 8 +- .../TypeTests/QdrantKeyTypeTests.cs | 2 +- .../TypeTests/RedisHashSetKeyTypeTests.cs | 8 +- .../TypeTests/RedisJsonKeyTypeTests.cs | 8 +- .../SqlServerCommandBuilderTests.cs | 76 ++---- .../TypeTests/SqlServerKeyTypeTests.cs | 4 +- .../TypeTests/SqliteKeyTypeTests.cs | 4 +- .../SqliteCommandBuilderTests.cs | 65 ++++- .../TypeTests/KeyTypeTests.cs | 247 +++++++++++++----- .../CollectionModelBuilderTests.cs | 10 +- 49 files changed, 1107 insertions(+), 623 deletions(-) diff --git a/dotnet/src/InternalUtilities/connectors/Memory/MongoDB/MongoModelBuilder.cs b/dotnet/src/InternalUtilities/connectors/Memory/MongoDB/MongoModelBuilder.cs index bb6d9e2daa2d..c8926258bd3c 100644 --- a/dotnet/src/InternalUtilities/connectors/Memory/MongoDB/MongoModelBuilder.cs +++ b/dotnet/src/InternalUtilities/connectors/Memory/MongoDB/MongoModelBuilder.cs @@ -42,11 +42,18 @@ protected override void ProcessTypeProperties(Type type, VectorStoreCollectionDe } } - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override bool SupportsKeyAutoGeneration(Type keyPropertyType) + => keyPropertyType == typeof(Guid) || keyPropertyType == typeof(ObjectId); + + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, int, long, Guid, ObjectId"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(int) || type == typeof(long) || type == typeof(Guid) || type == typeof(ObjectId); + if (type != typeof(string) && type != typeof(int) && type != typeof(long) && type != typeof(Guid) && type != typeof(ObjectId)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, int, long, Guid, ObjectId."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/AzureAISearch/AzureAISearchCollection.cs b/dotnet/src/VectorData/AzureAISearch/AzureAISearchCollection.cs index b82450e7a953..6d6af2ad0f0c 100644 --- a/dotnet/src/VectorData/AzureAISearch/AzureAISearchCollection.cs +++ b/dotnet/src/VectorData/AzureAISearch/AzureAISearchCollection.cs @@ -560,7 +560,19 @@ private async Task> MapToStorageModelAndUploadDoc (records, var generatedEmbeddings) = await ProcessEmbeddingsAsync(this._model, records, cancellationToken).ConfigureAwait(false); - var jsonObjects = records.Select((r, i) => this._mappper!.MapFromDataToStorageModel(r, i, generatedEmbeddings)); + // Handle auto-generated keys (client-side for Azure AI Search, which doesn't support server-side auto-generation) + var keyProperty = this._model.KeyProperty; + var jsonObjects = new List(); + var recordIndex = 0; + foreach (var record in records) + { + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + + jsonObjects.Add(this._mappper!.MapFromDataToStorageModel(record, recordIndex++, generatedEmbeddings)); + } return await this.RunOperationAsync( OperationName, diff --git a/dotnet/src/VectorData/AzureAISearch/AzureAISearchDynamicModelBuilder.cs b/dotnet/src/VectorData/AzureAISearch/AzureAISearchDynamicModelBuilder.cs index e51f11d32ae8..d0689c9b9089 100644 --- a/dotnet/src/VectorData/AzureAISearch/AzureAISearchDynamicModelBuilder.cs +++ b/dotnet/src/VectorData/AzureAISearch/AzureAISearchDynamicModelBuilder.cs @@ -17,8 +17,16 @@ internal class AzureAISearchDynamicModelBuilder() : CollectionModelBuilder(s_mod UsesExternalSerializer = true }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) - => AzureAISearchModelBuilder.IsKeyPropertyTypeValidCore(type, out supportedTypes); + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) + { + var type = keyProperty.Type; + + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } + } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) => AzureAISearchModelBuilder.IsDataPropertyTypeValidCore(type, out supportedTypes); diff --git a/dotnet/src/VectorData/AzureAISearch/AzureAISearchModelBuilder.cs b/dotnet/src/VectorData/AzureAISearch/AzureAISearchModelBuilder.cs index 0665942d6f63..31f90d77dd3c 100644 --- a/dotnet/src/VectorData/AzureAISearch/AzureAISearchModelBuilder.cs +++ b/dotnet/src/VectorData/AzureAISearch/AzureAISearchModelBuilder.cs @@ -21,8 +21,16 @@ internal class AzureAISearchModelBuilder() : CollectionJsonModelBuilder(s_modelB UsesExternalSerializer = true }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) - => IsKeyPropertyTypeValidCore(type, out supportedTypes); + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) + { + var type = keyProperty.Type; + + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } + } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) => IsDataPropertyTypeValidCore(type, out supportedTypes); @@ -30,13 +38,6 @@ protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] protected override bool IsVectorPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) => IsVectorPropertyTypeValidCore(type, out supportedTypes); - internal static bool IsKeyPropertyTypeValidCore(Type type, [NotNullWhen(false)] out string? supportedTypes) - { - supportedTypes = "string, Guid"; - - return type == typeof(string) || type == typeof(Guid); - } - internal static bool IsDataPropertyTypeValidCore(Type type, [NotNullWhen(false)] out string? supportedTypes) { supportedTypes = "string, int, long, double, float, bool, DateTimeOffset, or arrays/lists of these types"; diff --git a/dotnet/src/VectorData/CosmosMongoDB/CosmosMongoCollection.cs b/dotnet/src/VectorData/CosmosMongoDB/CosmosMongoCollection.cs index 6131f439f624..dcdf40209db8 100644 --- a/dotnet/src/VectorData/CosmosMongoDB/CosmosMongoCollection.cs +++ b/dotnet/src/VectorData/CosmosMongoDB/CosmosMongoCollection.cs @@ -254,6 +254,31 @@ private async Task UpsertCoreAsync(TRecord record, int recordIndex, IReadOnlyLis { const string OperationName = "ReplaceOne"; + // Handle auto-generated keys + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated) + { + switch (keyProperty.Type) + { + case var t when t == typeof(Guid): + if (keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + break; + + case var t when t == typeof(ObjectId): + if (keyProperty.GetValue(record) == ObjectId.Empty) + { + keyProperty.SetValue(record, ObjectId.GenerateNewId()); + } + break; + + default: + throw new UnreachableException(); + } + } + var replaceOptions = new ReplaceOptions { IsUpsert = true }; var storageModel = this._mapper.MapFromDataToStorageModel(record, recordIndex, generatedEmbeddings); diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs index 480ab596128d..8a99611db3f9 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs @@ -402,6 +402,13 @@ private async Task UpsertCoreAsync(TRecord record, int recordIndex, IReadOnlyLis { const string OperationName = "UpsertItem"; + // Handle auto-generated keys (client-side for Cosmos DB NoSQL, which doesn't support server-side auto-generation) + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + var jsonObject = this._mapper.MapFromDataToStorageModel(record, recordIndex, generatedEmbeddings); var keyValue = jsonObject.TryGetPropertyValue(this._model.KeyProperty.StorageName!, out var jsonKey) ? jsonKey?.ToString() : null; diff --git a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs index 6690a0ca150c..803d4aa4a2c9 100644 --- a/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs +++ b/dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlModelBuilder.cs @@ -21,11 +21,15 @@ internal class CosmosNoSqlModelBuilder() : CollectionJsonModelBuilder(s_modelBui ReservedKeyStorageName = CosmosNoSqlConstants.ReservedKeyPropertyName }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, Guid"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(Guid); + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/InMemory/InMemoryCollection.cs b/dotnet/src/VectorData/InMemory/InMemoryCollection.cs index 4144b453069e..5303d9abf66d 100644 --- a/dotnet/src/VectorData/InMemory/InMemoryCollection.cs +++ b/dotnet/src/VectorData/InMemory/InMemoryCollection.cs @@ -221,9 +221,17 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio var collectionDictionary = this.GetCollectionDictionary(); var recordIndex = 0; + var keyProperty = this._model.KeyProperty; foreach (var record in records) { - var key = (TKey)this._model.KeyProperty.GetValueAsObject(record)!; + var key = (TKey)keyProperty.GetValueAsObject(record)!; + if (keyProperty.IsAutoGenerated && (Guid)(object)key == Guid.Empty) + { + var generatedGuid = Guid.NewGuid(); + keyProperty.SetValue(record, generatedGuid); + key = (TKey)(object)generatedGuid; + } + var wrappedRecord = new InMemoryRecordWrapper(record); if (generatedEmbeddings is not null) diff --git a/dotnet/src/VectorData/InMemory/InMemoryModelBuilder.cs b/dotnet/src/VectorData/InMemory/InMemoryModelBuilder.cs index 1112b9a887f9..aab0cb1c5d10 100644 --- a/dotnet/src/VectorData/InMemory/InMemoryModelBuilder.cs +++ b/dotnet/src/VectorData/InMemory/InMemoryModelBuilder.cs @@ -18,12 +18,14 @@ internal class InMemoryModelBuilder() : CollectionModelBuilder(ValidationOptions SupportsMultipleVectors = true }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = ""; - - // All .NET types are supported by the InMemory provider - return true; + // All .NET types are supported by the InMemory provider, but we support auto-generation of keys only for GUIDs + if (keyProperty.IsAutoGenerated && keyProperty.Type != typeof(Guid)) + { + throw new NotSupportedException( + $"Auto-generation is only supported for key properties of type Guid. Property '{keyProperty.ModelName}' has type '{keyProperty.Type.Name}'."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/MongoDB/MongoCollection.cs b/dotnet/src/VectorData/MongoDB/MongoCollection.cs index c52d3c25bc6d..443e114b02d0 100644 --- a/dotnet/src/VectorData/MongoDB/MongoCollection.cs +++ b/dotnet/src/VectorData/MongoDB/MongoCollection.cs @@ -272,6 +272,32 @@ private async Task UpsertCoreAsync(TRecord record, int recordIndex, IReadOnlyLis { const string OperationName = "ReplaceOne"; + // Handle auto-generated keys + var keyProperty = this._model.KeyProperty; + + if (keyProperty.IsAutoGenerated) + { + switch (keyProperty.Type) + { + case var t when t == typeof(Guid): + if (keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + break; + + case var t when t == typeof(ObjectId): + if (keyProperty.GetValue(record) == ObjectId.Empty) + { + keyProperty.SetValue(record, ObjectId.GenerateNewId()); + } + break; + + default: + throw new UnreachableException(); + } + } + var replaceOptions = new ReplaceOptions { IsUpsert = true }; var storageModel = this._mapper.MapFromDataToStorageModel(record, recordIndex, generatedEmbeddings); diff --git a/dotnet/src/VectorData/PgVector/PostgresCollection.cs b/dotnet/src/VectorData/PgVector/PostgresCollection.cs index 1523c5babfea..9473de31afd0 100644 --- a/dotnet/src/VectorData/PgVector/PostgresCollection.cs +++ b/dotnet/src/VectorData/PgVector/PostgresCollection.cs @@ -166,69 +166,13 @@ public override async Task EnsureCollectionDeletedAsync(CancellationToken cancel } /// - public override async Task UpsertAsync(TRecord record, CancellationToken cancellationToken = default) - { - const string OperationName = "Upsert"; - - Dictionary>? generatedEmbeddings = null; - - var vectorPropertyCount = this._model.VectorProperties.Count; - for (var i = 0; i < vectorPropertyCount; i++) - { - var vectorProperty = this._model.VectorProperties[i]; - - if (PostgresModelBuilder.IsVectorPropertyTypeValidCore(vectorProperty.Type, out _)) - { - continue; - } - - // We have a vector property whose type isn't natively supported - we need to generate embeddings. - Debug.Assert(vectorProperty.EmbeddingGenerator is not null); - - // TODO: Ideally we'd group together vector properties using the same generator (and with the same input and output properties), - // and generate embeddings for them in a single batch. That's some more complexity though. - if (vectorProperty.TryGenerateEmbedding>(record, cancellationToken, out var floatTask)) - { - generatedEmbeddings ??= new Dictionary>(vectorPropertyCount); - generatedEmbeddings[vectorProperty] = [await floatTask.ConfigureAwait(false)]; - } -#if NET - else if (vectorProperty.TryGenerateEmbedding>(record, cancellationToken, out var halfTask)) - { - generatedEmbeddings ??= new Dictionary>(vectorPropertyCount); - generatedEmbeddings[vectorProperty] = [await halfTask.ConfigureAwait(false)]; - } -#endif - else if (vectorProperty.TryGenerateEmbedding(record, cancellationToken, out var binaryTask)) - { - generatedEmbeddings ??= new Dictionary>(vectorPropertyCount); - generatedEmbeddings[vectorProperty] = [await binaryTask.ConfigureAwait(false)]; - } - else - { - throw new InvalidOperationException( - $"The embedding generator configured on property '{vectorProperty.ModelName}' cannot produce an embedding of type '{typeof(Embedding).Name}' for the given input type."); - } - } - - using var connection = await this._dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false); - using var command = connection.CreateCommand(); - - // using var command = PostgresSqlBuilder.BuildUpsertCommand(this._schema, this.Name, this._model, record, generatedEmbeddings); - // command.Connection = connection; - _ = PostgresSqlBuilder.BuildUpsertCommand(command, this._schema, this.Name, this._model, [record], generatedEmbeddings); - - await this.RunOperationAsync(OperationName, () => command.ExecuteNonQueryAsync(cancellationToken)) - .ConfigureAwait(false); - } + public override Task UpsertAsync(TRecord record, CancellationToken cancellationToken = default) + => this.UpsertAsync([record], cancellationToken); /// public override async Task UpsertAsync(IEnumerable records, CancellationToken cancellationToken = default) { Verify.NotNull(records); - - const string OperationName = "UpsertBatch"; - IReadOnlyList? recordsList = null; // If an embedding generator is defined, invoke it once per property for all records. @@ -274,6 +218,11 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio generatedEmbeddings[vectorProperty] = await halfTask.ConfigureAwait(false); } #endif + else if (vectorProperty.TryGenerateEmbeddings(records, cancellationToken, out var binaryTask)) + { + generatedEmbeddings ??= new Dictionary>(vectorPropertyCount); + generatedEmbeddings[vectorProperty] = await binaryTask.ConfigureAwait(false); + } else { throw new InvalidOperationException( @@ -282,11 +231,54 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio } using var connection = await this._dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false); - using var command = connection.CreateCommand(); + using var batch = connection.CreateBatch(); - if (PostgresSqlBuilder.BuildUpsertCommand(command, this._schema, this.Name, this._model, records, generatedEmbeddings)) + // If key auto-generation is enabled, we'll need to enumerate over the records multiple times: + // once to generate the upsert batch, and again to populate the retrieved keys into the records. + // To prevent multiple enumeration of the input enumerable, we materialize it here if needed. + if (this._model.KeyProperty.IsAutoGenerated && recordsList is null) { - await this.RunOperationAsync(OperationName, () => command.ExecuteNonQueryAsync(cancellationToken)).ConfigureAwait(false); + recordsList = records is IReadOnlyList r ? r : records.ToList(); + + if (recordsList.Count == 0) + { + return; + } + + records = recordsList; + } + + if (PostgresSqlBuilder.BuildUpsertCommand(batch, this._schema, this.Name, this._model, records, generatedEmbeddings)) + { + await this.RunOperationAsync("Upsert", async () => + { + var reader = await batch.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false); + try + { + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated) + { + int? keyOrdinal = null; + + foreach (var record in recordsList!) + { + if (Equals(keyProperty.GetValueAsObject(record), default(TKey))) + { + keyOrdinal ??= reader.GetOrdinal(keyProperty.StorageName); + + await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + var keyValue = reader.GetFieldValue(keyOrdinal.Value); + this._model.KeyProperty.SetValue(record, keyValue); + await reader.NextResultAsync(cancellationToken).ConfigureAwait(false); + } + } + } + } + finally + { + await reader.DisposeAsync().ConfigureAwait(false); + } + }).ConfigureAwait(false); } } @@ -513,11 +505,13 @@ private async Task InternalCreateCollectionAsync(bool ifNotExists, CancellationT await using (connection) { + var pgVersion = connection.PostgreSqlVersion; + // Prepare the SQL commands. using var batch = connection.CreateBatch(); batch.BatchCommands.Add( - new NpgsqlBatchCommand(PostgresSqlBuilder.BuildCreateTableSql(this._schema, this.Name, this._model, ifNotExists))); + new NpgsqlBatchCommand(PostgresSqlBuilder.BuildCreateTableSql(this._schema, this.Name, this._model, pgVersion, ifNotExists))); foreach (var (column, kind, function, isVector) in PostgresPropertyMapping.GetIndexInfo(this._model.Properties)) { diff --git a/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs b/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs index d52b9c53e1af..e6e54fdac8fa 100644 --- a/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs +++ b/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs @@ -21,15 +21,22 @@ internal class PostgresModelBuilder() : CollectionModelBuilder(PostgresModelBuil SupportsMultipleVectors = true, }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override bool SupportsKeyAutoGeneration(Type keyPropertyType) + => keyPropertyType == typeof(Guid) || keyPropertyType == typeof(int) || keyPropertyType == typeof(long); + + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "short, int, long, string, Guid"; + var type = keyProperty.Type; - return type == typeof(short) - || type == typeof(int) - || type == typeof(long) - || type == typeof(string) - || type == typeof(Guid); + if (type != typeof(short) + && type != typeof(int) + && type != typeof(long) + && type != typeof(string) + && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: short, int, long, string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/PgVector/PostgresSqlBuilder.cs b/dotnet/src/VectorData/PgVector/PostgresSqlBuilder.cs index 16deaa8bd0b5..8282e6f7ec0c 100644 --- a/dotnet/src/VectorData/PgVector/PostgresSqlBuilder.cs +++ b/dotnet/src/VectorData/PgVector/PostgresSqlBuilder.cs @@ -46,7 +46,7 @@ FROM information_schema.tables command.Parameters.Add(new() { Value = schema }); } - internal static string BuildCreateTableSql(string schema, string tableName, CollectionModel model, bool ifNotExists = true) + internal static string BuildCreateTableSql(string schema, string tableName, CollectionModel model, Version pgVersion, bool ifNotExists = true) { if (string.IsNullOrWhiteSpace(tableName)) { @@ -64,11 +64,26 @@ internal static string BuildCreateTableSql(string schema, string tableName, Coll createTableCommand.AppendIdentifier(schema).Append('.').AppendIdentifier(tableName).AppendLine(" ("); // Add the key column - var keyPgTypeInfo = PostgresPropertyMapping.GetPostgresTypeName(model.KeyProperty.Type); - createTableCommand.Append(" ").AppendIdentifier(keyName).Append(' ').Append(keyPgTypeInfo.PgType); - if (!keyPgTypeInfo.IsNullable) + var keyStoreType = PostgresPropertyMapping.GetPostgresTypeName(model.KeyProperty.Type).PgType; + createTableCommand.Append(" ").AppendIdentifier(keyName).Append(' ').Append(keyStoreType); + if (model.KeyProperty.IsAutoGenerated) { - createTableCommand.Append(" NOT NULL"); + switch (keyStoreType.ToUpperInvariant()) + { + case "INTEGER": + case "BIGINT": + createTableCommand.Append(" GENERATED BY DEFAULT AS IDENTITY"); + break; + case "UUID": + // UUIDv7 is superior for PostgreSQL indexing, but generating it was only added in PG18. + // Fall back to UUIDv4 in older PG versions. + createTableCommand.Append(pgVersion >= new Version(18, 0) + ? " DEFAULT uuidv7()" + : " DEFAULT gen_random_uuid()"); + break; + default: + throw new NotSupportedException($"Auto-generation of keys is not supported for key type '{keyStoreType}'. Only INTEGER, BIGINT and UUID are supported."); + } } createTableCommand.AppendLine(","); @@ -157,53 +172,35 @@ internal static void BuildDropTableCommand(NpgsqlCommand command, string schema, } /// - internal static bool BuildUpsertCommand( - NpgsqlCommand command, + internal static bool BuildUpsertCommand( + NpgsqlBatch batch, string schema, string tableName, CollectionModel model, IEnumerable records, Dictionary>? generatedEmbeddings) { - StringBuilder sql = new(); - - sql.Append("INSERT INTO ").AppendIdentifier(schema).Append('.').AppendIdentifier(tableName).Append(" ("); - - for (var i = 0; i < model.Properties.Count; i++) - { - var property = model.Properties[i]; - - if (i > 0) - { - sql.Append(", "); - } - - sql.AppendIdentifier(property.StorageName); - } - - sql - .AppendLine(")") - .Append("VALUES "); + // Note: since keys may be auto-generated, we can't use a single multi-value INSERT statement, since that would return + // the generated keys in random order. Use a batch of single-value INSERT statements instead. + string? upsertSql = null, upsertSqlWithAutoGeneratedKey = null; + var keyProperty = model.KeyProperty; var recordIndex = 0; - var parameterIndex = 1; foreach (var record in records) { - if (recordIndex > 0) - { - sql.Append(", "); - } + var batchCommand = batch.CreateBatchCommand(); - sql.Append('('); + var autoGeneratedKey = keyProperty.IsAutoGenerated && Equals(keyProperty.GetValueAsObject(record), default(TKey)); + batchCommand.CommandText = autoGeneratedKey + ? upsertSqlWithAutoGeneratedKey ??= GenerateSingleUpsertSql(autoGeneratedKey: true) + : upsertSql ??= GenerateSingleUpsertSql(autoGeneratedKey: false); - for (var i = 0; i < model.Properties.Count; i++) + foreach (var property in model.Properties) { - var property = model.Properties[i]; - - if (i > 0) + if (property is KeyPropertyModel && autoGeneratedKey) { - sql.Append(", "); + continue; } var value = property.GetValueAsObject(record); @@ -218,12 +215,10 @@ internal static bool BuildUpsertCommand( value = PostgresPropertyMapping.MapVectorForStorageModel(value); } - command.Parameters.Add(new() { Value = value ?? DBNull.Value }); - sql.Append('$').Append(parameterIndex++); + batchCommand.Parameters.Add(new() { Value = value ?? DBNull.Value }); } - sql.Append(')'); - + batch.BatchCommands.Add(batchCommand); recordIndex++; } @@ -233,33 +228,96 @@ internal static bool BuildUpsertCommand( return false; } - sql - .AppendLine() - .Append("ON CONFLICT (").AppendIdentifier(model.KeyProperty.StorageName).Append(')'); - - sql - .AppendLine() - .AppendLine("DO UPDATE SET "); + return true; - var propertyIndex = 0; - foreach (var property in model.Properties) + string GenerateSingleUpsertSql(bool autoGeneratedKey) { - if (property is KeyPropertyModel) + StringBuilder sqlBuilder = new(); + + sqlBuilder + .Append("INSERT INTO ") + .AppendIdentifier(schema) + .Append('.') + .AppendIdentifier(tableName) + .Append(" ("); + + var i = 0; + foreach (var property in model.Properties) { - continue; + if (property is KeyPropertyModel && autoGeneratedKey) + { + continue; + } + + if (i++ > 0) + { + sqlBuilder.Append(", "); + } + + sqlBuilder.AppendIdentifier(property.StorageName); } - if (propertyIndex++ > 0) + sqlBuilder + .AppendLine(")") + .Append("VALUES ("); + + i = 0; + foreach (var property in model.Properties) { - sql.AppendLine(", "); + if (property is KeyPropertyModel && autoGeneratedKey) + { + continue; + } + + if (i++ > 0) + { + sqlBuilder.Append(", "); + } + + sqlBuilder.Append('$').Append(i); } - sql.Append(" ").AppendIdentifier(property.StorageName).Append(" = EXCLUDED.").AppendIdentifier(property.StorageName); - } + sqlBuilder.Append(')'); - command.CommandText = sql.ToString(); + if (autoGeneratedKey) + { + sqlBuilder + .AppendLine() + .Append("RETURNING ") + .AppendIdentifier(keyProperty.StorageName); + } + else + { + sqlBuilder + .AppendLine() + .Append("ON CONFLICT (") + .AppendIdentifier(keyProperty.StorageName) + .AppendLine(")") + .AppendLine("DO UPDATE SET "); + + i = 0; + foreach (var property in model.Properties) + { + if (property is KeyPropertyModel) + { + continue; + } - return true; + if (i++ > 0) + { + sqlBuilder.AppendLine(", "); + } + + sqlBuilder + .Append(" ") + .AppendIdentifier(property.StorageName) + .Append(" = EXCLUDED.") + .AppendIdentifier(property.StorageName); + } + } + + return sqlBuilder.ToString(); + } } /// diff --git a/dotnet/src/VectorData/Pinecone/PineconeCollection.cs b/dotnet/src/VectorData/Pinecone/PineconeCollection.cs index f6e8fac90bd5..d545c6c2a887 100644 --- a/dotnet/src/VectorData/Pinecone/PineconeCollection.cs +++ b/dotnet/src/VectorData/Pinecone/PineconeCollection.cs @@ -314,6 +314,13 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell { Verify.NotNull(record); + // Handle auto-generated keys (client-side for Pinecone, which doesn't support server-side auto-generation) + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + // If an embedding generator is defined, invoke it once for all records. Embedding? generatedEmbedding = null; @@ -382,7 +389,19 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio } } - var vectors = records.Select((r, i) => this._mapper.MapFromDataToStorageModel(r, generatedEmbeddings?[i])).ToList(); + // Handle auto-generated keys (client-side for Pinecone, which doesn't support server-side auto-generation) + var keyProperty = this._model.KeyProperty; + var vectors = new List(); + var recordIndex = 0; + foreach (var record in records) + { + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + + vectors.Add(this._mapper.MapFromDataToStorageModel(record, generatedEmbeddings?[recordIndex++])); + } if (vectors.Count == 0) { diff --git a/dotnet/src/VectorData/Pinecone/PineconeModelBuilder.cs b/dotnet/src/VectorData/Pinecone/PineconeModelBuilder.cs index cdba6c103e92..e71b9effd0f0 100644 --- a/dotnet/src/VectorData/Pinecone/PineconeModelBuilder.cs +++ b/dotnet/src/VectorData/Pinecone/PineconeModelBuilder.cs @@ -19,11 +19,15 @@ internal class PineconeModelBuilder() : CollectionModelBuilder(s_validationOptio SupportsMultipleVectors = false, }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, Guid"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(Guid); + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/Qdrant/QdrantCollection.cs b/dotnet/src/VectorData/Qdrant/QdrantCollection.cs index 2dfaf498a869..0ded34a73330 100644 --- a/dotnet/src/VectorData/Qdrant/QdrantCollection.cs +++ b/dotnet/src/VectorData/Qdrant/QdrantCollection.cs @@ -504,9 +504,20 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio } // Create points from records. - var pointStructs = records.Select((r, i) => this._mapper.MapFromDataToStorageModel(r, i, generatedEmbeddings)).ToList(); + var keyProperty = this._model.KeyProperty; + var pointStructs = new List(); + var recordIndex = 0; + foreach (var record in records) + { + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + + pointStructs.Add(this._mapper.MapFromDataToStorageModel(record, recordIndex++, generatedEmbeddings)); + } - if (pointStructs is { Count: 0 }) + if (pointStructs.Count == 0) { return; } diff --git a/dotnet/src/VectorData/Qdrant/QdrantModelBuilder.cs b/dotnet/src/VectorData/Qdrant/QdrantModelBuilder.cs index 57d4395c6289..a290c303b85e 100644 --- a/dotnet/src/VectorData/Qdrant/QdrantModelBuilder.cs +++ b/dotnet/src/VectorData/Qdrant/QdrantModelBuilder.cs @@ -20,11 +20,15 @@ private static CollectionModelBuildingOptions GetModelBuildOptions(bool hasNamed SupportsMultipleVectors = hasNamedVectors, }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "ulong, Guid"; + var type = keyProperty.Type; - return type == typeof(ulong) || type == typeof(Guid); + if (type != typeof(ulong) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be either ulong or Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/Redis/RedisHashSetCollection.cs b/dotnet/src/VectorData/Redis/RedisHashSetCollection.cs index 34dd631df5f8..7dc66f8880f7 100644 --- a/dotnet/src/VectorData/Redis/RedisHashSetCollection.cs +++ b/dotnet/src/VectorData/Redis/RedisHashSetCollection.cs @@ -305,6 +305,13 @@ private async Task UpsertCoreAsync(TRecord record, int recordIndex, IReadOnlyLis { Verify.NotNull(record); + // Auto-generate key if needed (client-side for Redis) + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + // Map. var redisHashSetRecord = this._mapper.MapFromDataToStorageModel(record, recordIndex, generatedEmbeddings); diff --git a/dotnet/src/VectorData/Redis/RedisJsonCollection.cs b/dotnet/src/VectorData/Redis/RedisJsonCollection.cs index d570b4473e1e..ebed8c5bb6a5 100644 --- a/dotnet/src/VectorData/Redis/RedisJsonCollection.cs +++ b/dotnet/src/VectorData/Redis/RedisJsonCollection.cs @@ -357,6 +357,13 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell { Verify.NotNull(record); + // Auto-generate key if needed (client-side for Redis) + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + // Map. (_, var generatedEmbeddings) = await RedisFieldMapping.ProcessEmbeddingsAsync(this._model, [record], cancellationToken).ConfigureAwait(false); @@ -385,11 +392,18 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio (records, var generatedEmbeddings) = await RedisFieldMapping.ProcessEmbeddingsAsync(this._model, records, cancellationToken).ConfigureAwait(false); var redisRecords = new List<(string maybePrefixedKey, string originalKey, string serializedRecord)>(); + var keyProperty = this._model.KeyProperty; var recordIndex = 0; foreach (var record in records) { + // Auto-generate key if needed (client-side for Redis) + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + var mapResult = this._mapper.MapFromDataToStorageModel(record, recordIndex++, generatedEmbeddings); var serializedRecord = JsonSerializer.Serialize(mapResult.Node, this._jsonSerializerOptions); var redisJsonRecord = new { Key = mapResult.Key, SerializedRecord = serializedRecord }; diff --git a/dotnet/src/VectorData/Redis/RedisJsonDynamicModelBuilder.cs b/dotnet/src/VectorData/Redis/RedisJsonDynamicModelBuilder.cs index 5edc3e68d1ce..ffebe5637cba 100644 --- a/dotnet/src/VectorData/Redis/RedisJsonDynamicModelBuilder.cs +++ b/dotnet/src/VectorData/Redis/RedisJsonDynamicModelBuilder.cs @@ -17,11 +17,15 @@ internal class RedisJsonDynamicModelBuilder(CollectionModelBuildingOptions optio => vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType) ?? vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType); - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, Guid"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(Guid); + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/Redis/RedisJsonModelBuilder.cs b/dotnet/src/VectorData/Redis/RedisJsonModelBuilder.cs index 28a1ac7da516..44b0cbd996af 100644 --- a/dotnet/src/VectorData/Redis/RedisJsonModelBuilder.cs +++ b/dotnet/src/VectorData/Redis/RedisJsonModelBuilder.cs @@ -18,11 +18,15 @@ internal class RedisJsonModelBuilder(CollectionModelBuildingOptions options) : C => vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType) ?? vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType); - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, Guid"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(Guid); + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/Redis/RedisModelBuilder.cs b/dotnet/src/VectorData/Redis/RedisModelBuilder.cs index 328446ed13f8..47ff2d3e6b67 100644 --- a/dotnet/src/VectorData/Redis/RedisModelBuilder.cs +++ b/dotnet/src/VectorData/Redis/RedisModelBuilder.cs @@ -20,11 +20,15 @@ internal class RedisModelBuilder(CollectionModelBuildingOptions options) : Colle => vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType) ?? vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType); - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, Guid"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(Guid); + if (type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/SqlServer/SqlServerCollection.cs b/dotnet/src/VectorData/SqlServer/SqlServerCollection.cs index b9026f442209..f7695ec00fde 100644 --- a/dotnet/src/VectorData/SqlServer/SqlServerCollection.cs +++ b/dotnet/src/VectorData/SqlServer/SqlServerCollection.cs @@ -374,14 +374,18 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell } using SqlConnection connection = new(this._connectionString); - using SqlCommand command = SqlServerCommandBuilder.MergeIntoSingle( - connection, + using SqlCommand command = connection.CreateCommand(); + SqlServerCommandBuilder.Upsert( + command, this._schema, this.Name, this._model, - record, + [record], + firstRecordIndex: 0, generatedEmbeddings); + var keyProperty = this._model.KeyProperty; + await connection.ExecuteWithErrorHandlingAsync( this._collectionMetadata, "Upsert", @@ -389,8 +393,15 @@ await connection.ExecuteWithErrorHandlingAsync( { using SqlDataReader reader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false); await reader.ReadAsync(cancellationToken).ConfigureAwait(false); - // TODO: Currently unused (#11835), but will be injected into the record in the future. - return reader.GetFieldValue(0); + + // Inject the generated key into the record if auto-generation was used + if (keyProperty.IsAutoGenerated && Equals(keyProperty.GetValueAsObject(record), default(TKey))) + { + var keyValue = reader.GetFieldValue(0); + keyProperty.SetValue(record, keyValue); + } + + return 0; }, cancellationToken).ConfigureAwait(false); } @@ -446,17 +457,41 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio } } + // If key auto-generation is enabled, we need to read back generated keys and inject them into records. + // Materialize the records' enumerable if needed, to allow iteration for key injection. + var keyProperty = this._model.KeyProperty; + if (keyProperty.IsAutoGenerated && recordsList is null) + { + recordsList = records is IReadOnlyList r ? r : records.ToList(); + + if (recordsList.Count == 0) + { + return; + } + + records = recordsList; + } + using SqlConnection connection = new(this._connectionString); await connection.OpenAsync(cancellationToken).ConfigureAwait(false); using SqlTransaction transaction = connection.BeginTransaction(); int parametersPerRecord = this._model.Properties.Count; int taken = 0; + int batchSize = SqlServerConstants.MaxParameterCount / parametersPerRecord; try { while (true) { + // Materialize the batch to a list so we can iterate multiple times: + // once for building the command, once for reading back results. + var batch = records.Skip(taken).Take(batchSize).ToList(); + if (batch.Count == 0) + { + break; + } + #if NET SqlCommand command = new("", connection, transaction); await using (command.ConfigureAwait(false)) @@ -464,24 +499,41 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio using (SqlCommand command = new("", connection, transaction)) #endif { - if (!SqlServerCommandBuilder.MergeIntoMany( + if (!SqlServerCommandBuilder.Upsert( command, this._schema, this.Name, this._model, - records.Skip(taken).Take(SqlServerConstants.MaxParameterCount / parametersPerRecord), + batch, firstRecordIndex: taken, generatedEmbeddings)) { - break; // records is empty + break; // records is empty (shouldn't happen given check above, but defensive) } - checked + // Execute and read back the generated keys. + // Each MERGE statement returns a single result set with one row containing the key. + using SqlDataReader reader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false); + + // Iterate through the records in this batch and inject generated keys where needed. + foreach (var record in batch) { - taken += (command.Parameters.Count / parametersPerRecord); + await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + + // Only inject key if auto-generation is enabled and record had a default key value + if (keyProperty.IsAutoGenerated && Equals(keyProperty.GetValueAsObject(record), default(TKey))) + { + var keyValue = reader.GetFieldValue(0); + keyProperty.SetValue(record, keyValue); + } + + await reader.NextResultAsync(cancellationToken).ConfigureAwait(false); } - await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false); + checked + { + taken += batch.Count; + } } } diff --git a/dotnet/src/VectorData/SqlServer/SqlServerCommandBuilder.cs b/dotnet/src/VectorData/SqlServer/SqlServerCommandBuilder.cs index 0879737ebcc5..139bcd91ed4c 100644 --- a/dotnet/src/VectorData/SqlServer/SqlServerCommandBuilder.cs +++ b/dotnet/src/VectorData/SqlServer/SqlServerCommandBuilder.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics; using System.Linq.Expressions; using System.Text; using System.Text.Json; @@ -35,7 +36,27 @@ internal static SqlCommand CreateTable( sb.Append("CREATE TABLE "); sb.AppendTableName(schema, tableName); sb.AppendLine(" ("); - sb.AppendIdentifier(model.KeyProperty.StorageName).Append(' ').Append(Map(model.KeyProperty)).AppendLine(" NOT NULL,"); + + var keyStoreType = Map(model.KeyProperty); + sb.AppendIdentifier(model.KeyProperty.StorageName).Append(' ').Append(keyStoreType); + if (model.KeyProperty.IsAutoGenerated) + { + switch (keyStoreType.ToUpperInvariant()) + { + case "SMALLINT": + case "INT": + case "BIGINT": + sb.Append(" IDENTITY"); + break; + case "UNIQUEIDENTIFIER": + sb.Append(" DEFAULT NEWSEQUENTIALID()"); + break; + default: + throw new UnreachableException(); + } + } + + sb.AppendLine(","); foreach (var property in model.DataProperties) { @@ -122,65 +143,25 @@ FROM INFORMATION_SCHEMA.TABLES return command; } - internal static SqlCommand MergeIntoSingle( - SqlConnection connection, - string? schema, - string tableName, - CollectionModel model, - object record, - Dictionary>? generatedEmbeddings) + /// + /// Checks if the key property uses SQL Server IDENTITY (for int/bigint) as opposed to DEFAULT (for GUID). + /// IDENTITY columns require SET IDENTITY_INSERT ON to insert explicit values. + /// + private static bool UsesIdentity(KeyPropertyModel keyProperty) { - SqlCommand command = connection.CreateCommand(); - StringBuilder sb = new(200); - sb.Append("MERGE INTO "); - sb.AppendTableName(schema, tableName); - sb.AppendLine(" AS t"); - sb.Append("USING (VALUES ("); - int paramIndex = 0; - - foreach (var property in model.Properties) + if (!keyProperty.IsAutoGenerated) { - sb.AppendParameterName(property, ref paramIndex, out var paramName).Append(','); - - var value = property is VectorPropertyModel vectorProperty && generatedEmbeddings?.TryGetValue(vectorProperty, out var ge) == true - ? ge[0] - : property.GetValueAsObject(record); - - command.AddParameter(property, paramName, value); - } - - sb[sb.Length - 1] = ')'; // replace the last comma with a closing parenthesis - sb.Append(") AS s ("); - sb.AppendIdentifiers(model.Properties); - sb.AppendLine(")"); - sb.Append("ON (t.").AppendIdentifier(model.KeyProperty.StorageName).Append(" = s.").AppendIdentifier(model.KeyProperty.StorageName).AppendLine(")"); - sb.AppendLine("WHEN MATCHED THEN"); - sb.Append("UPDATE SET "); - foreach (var property in model.Properties) - { - if (property is not KeyPropertyModel) // don't update the key - { - sb.Append("t.").AppendIdentifier(property.StorageName).Append(" = s.").AppendIdentifier(property.StorageName).Append(','); - } + return false; } - --sb.Length; // remove the last comma - sb.AppendLine(); - - sb.Append("WHEN NOT MATCHED THEN"); - sb.AppendLine(); - sb.Append("INSERT ("); - sb.AppendIdentifiers(model.Properties); - sb.AppendLine(")"); - sb.Append("VALUES ("); - sb.AppendIdentifiers(model.Properties, prefix: "s."); - sb.AppendLine(")"); - sb.Append("OUTPUT inserted.").AppendIdentifier(model.KeyProperty.StorageName).Append(';'); - command.CommandText = sb.ToString(); - return command; + var keyStoreType = Map(keyProperty).ToUpperInvariant(); + return keyStoreType is "SMALLINT" or "INT" or "BIGINT"; } - internal static bool MergeIntoMany( + // Note: since keys may be auto-generated, we can't use a single multi-value MERGE statement, since that would return + // the generated keys in undefined order (OUTPUT order is not guaranteed in MERGE). + // Use a batch of single-row MERGE statements instead - each returns a separate result set. + internal static bool Upsert( SqlCommand command, string? schema, string tableName, @@ -189,22 +170,42 @@ internal static bool MergeIntoMany( int firstRecordIndex, Dictionary>? generatedEmbeddings) { - StringBuilder sb = new(200); - // The DECLARE statement creates a table variable to store the keys of the inserted rows. - sb.AppendFormat("DECLARE @InsertedKeys TABLE (KeyColumn {0});", Map(model.KeyProperty)); - sb.AppendLine(); - // The MERGE statement performs the upsert operation and outputs the keys of the inserted rows into the table variable. - sb.Append("MERGE INTO "); - sb.AppendTableName(schema, tableName); - sb.AppendLine(" AS t"); // t stands for target - sb.AppendLine("USING (VALUES"); + var keyProperty = model.KeyProperty; + StringBuilder sb = new(500); + int rowIndex = 0, paramIndex = 0; + foreach (var record in records) { - sb.Append('('); + // A record needs auto-generation if the key property is auto-generated AND the record has a default key value. + var needsKeyGeneration = keyProperty.IsAutoGenerated && Equals(keyProperty.GetValueAsObject(record), default(TKey)); + // Skip key in INSERT when auto-generating (IDENTITY will provide the value) + var skipKeyInInsert = needsKeyGeneration; + // For explicit keys with IDENTITY columns, we need to enable IDENTITY_INSERT + // (only for int/bigint, not for GUID which uses DEFAULT NEWSEQUENTIALID()) + var needsIdentityInsert = UsesIdentity(keyProperty) && !needsKeyGeneration; + + // Enable IDENTITY_INSERT if we're inserting an explicit value into an IDENTITY column + if (needsIdentityInsert) + { + sb.Append("SET IDENTITY_INSERT "); + sb.AppendTableName(schema, tableName); + sb.AppendLine(" ON;"); + } + + sb.Append("MERGE INTO "); + sb.AppendTableName(schema, tableName); + sb.AppendLine(" AS t"); + sb.Append("USING (VALUES ("); foreach (var property in model.Properties) { + // Skip key in VALUES when auto-generating + if (property is KeyPropertyModel && skipKeyInInsert) + { + continue; + } + sb.AppendParameterName(property, ref paramIndex, out var paramName).Append(','); var value = property is VectorPropertyModel vectorProperty && generatedEmbeddings?.TryGetValue(vectorProperty, out var ge) == true @@ -215,7 +216,51 @@ internal static bool MergeIntoMany( } sb[sb.Length - 1] = ')'; // replace the last comma with a closing parenthesis - sb.AppendLine(","); + sb.Append(") AS s ("); + sb.AppendIdentifiers(model.Properties, skipKey: skipKeyInInsert); + sb.AppendLine(")"); + + if (needsKeyGeneration) + { + // When auto-generating a key, we always insert (ON condition never matches). + sb.AppendLine("ON (1=0)"); + } + else + { + // For upsert, match on the key from the source + sb.Append("ON (t.").AppendIdentifier(model.KeyProperty.StorageName).Append(" = s.").AppendIdentifier(model.KeyProperty.StorageName).AppendLine(")"); + sb.AppendLine("WHEN MATCHED THEN"); + sb.Append("UPDATE SET "); + foreach (var property in model.Properties) + { + if (property is not KeyPropertyModel) // don't update the key + { + sb.Append("t.").AppendIdentifier(property.StorageName).Append(" = s.").AppendIdentifier(property.StorageName).Append(','); + } + } + --sb.Length; // remove the last comma + sb.AppendLine(); + } + + sb.AppendLine("WHEN NOT MATCHED THEN"); + sb.Append("INSERT ("); + sb.AppendIdentifiers(model.Properties, skipKey: skipKeyInInsert); + sb.AppendLine(")"); + sb.Append("VALUES ("); + sb.AppendIdentifiers(model.Properties, prefix: "s.", skipKey: skipKeyInInsert); + sb.AppendLine(")"); + sb.Append("OUTPUT inserted.").AppendIdentifier(model.KeyProperty.StorageName).AppendLine(";"); + + // Disable IDENTITY_INSERT after the MERGE + if (needsIdentityInsert) + { + sb.Append("SET IDENTITY_INSERT "); + sb.AppendTableName(schema, tableName); + sb.AppendLine(" OFF;"); + } + + sb.AppendLine(); + rowIndex++; } @@ -224,36 +269,6 @@ internal static bool MergeIntoMany( return false; // there is nothing to do! } - sb.Length -= (1 + Environment.NewLine.Length); // remove the last comma and newline - - sb.Append(") AS s ("); // s stands for source - sb.AppendIdentifiers(model.Properties); - sb.AppendLine(")"); - sb.Append("ON (t.").AppendIdentifier(model.KeyProperty.StorageName).Append(" = s.").AppendIdentifier(model.KeyProperty.StorageName).AppendLine(")"); - sb.AppendLine("WHEN MATCHED THEN"); - sb.Append("UPDATE SET "); - foreach (var property in model.Properties) - { - if (property is not KeyPropertyModel) // don't update the key - { - sb.Append("t.").AppendIdentifier(property.StorageName).Append(" = s.").AppendIdentifier(property.StorageName).Append(','); - } - } - --sb.Length; // remove the last comma - sb.AppendLine(); - sb.Append("WHEN NOT MATCHED THEN"); - sb.AppendLine(); - sb.Append("INSERT ("); - sb.AppendIdentifiers(model.Properties); - sb.AppendLine(")"); - sb.Append("VALUES ("); - sb.AppendIdentifiers(model.Properties, prefix: "s."); - sb.AppendLine(")"); - sb.Append("OUTPUT inserted.").AppendIdentifier(model.KeyProperty.StorageName).AppendLine(" INTO @InsertedKeys (KeyColumn);"); - - // The SELECT statement returns the keys of the inserted rows. - sb.Append("SELECT KeyColumn FROM @InsertedKeys;"); - command.CommandText = sb.ToString(); return true; } @@ -518,7 +533,8 @@ internal static StringBuilder AppendIdentifier(this StringBuilder sb, string ide private static StringBuilder AppendIdentifiers(this StringBuilder sb, IEnumerable properties, string? prefix = null, - bool includeVectors = true) + bool includeVectors = true, + bool skipKey = false) { bool any = false; foreach (var property in properties) @@ -528,6 +544,11 @@ private static StringBuilder AppendIdentifiers(this StringBuilder sb, continue; } + if (skipKey && property is KeyPropertyModel) + { + continue; + } + if (prefix is not null) { sb.Append(prefix); diff --git a/dotnet/src/VectorData/SqlServer/SqlServerModelBuilder.cs b/dotnet/src/VectorData/SqlServer/SqlServerModelBuilder.cs index 70e444e88544..c72513e4f913 100644 --- a/dotnet/src/VectorData/SqlServer/SqlServerModelBuilder.cs +++ b/dotnet/src/VectorData/SqlServer/SqlServerModelBuilder.cs @@ -20,14 +20,18 @@ internal class SqlServerModelBuilder() : CollectionModelBuilder(s_modelBuildingO SupportsMultipleVectors = true, }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override bool SupportsKeyAutoGeneration(Type keyPropertyType) + => keyPropertyType == typeof(Guid) || keyPropertyType == typeof(int) || keyPropertyType == typeof(long); + + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "int, long, string, Guid"; + var type = keyProperty.Type; - return type == typeof(int) // INT - || type == typeof(long) // BIGINT - || type == typeof(string) // VARCHAR - || type == typeof(Guid); // UNIQUEIDENTIFIER + if (type != typeof(int) && type != typeof(long) && type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: int, long, string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/SqliteVec/SqliteCollection.cs b/dotnet/src/VectorData/SqliteVec/SqliteCollection.cs index 71529dc67aa2..a1b404d6ece2 100644 --- a/dotnet/src/VectorData/SqliteVec/SqliteCollection.cs +++ b/dotnet/src/VectorData/SqliteVec/SqliteCollection.cs @@ -349,98 +349,15 @@ public override async IAsyncEnumerable GetAsync(IEnumerable keys, } /// - public override async Task UpsertAsync(TRecord record, CancellationToken cancellationToken = default) - { - Verify.NotNull(record); - - Dictionary>>? generatedEmbeddings = null; - - var vectorPropertyCount = this._model.VectorProperties.Count; - for (var i = 0; i < vectorPropertyCount; i++) - { - var vectorProperty = this._model.VectorProperties[i]; - - if (SqliteModelBuilder.IsVectorPropertyTypeValidCore(vectorProperty.Type, out _)) - { - continue; - } - - // We have a vector property whose type isn't natively supported - we need to generate embeddings. - Debug.Assert(vectorProperty.EmbeddingGenerator is not null); - - // TODO: Ideally we'd group together vector properties using the same generator (and with the same input and output properties), - // and generate embeddings for them in a single batch. That's some more complexity though. - if (vectorProperty.TryGenerateEmbedding>(record, cancellationToken, out var floatTask)) - { - generatedEmbeddings ??= new Dictionary>>(vectorPropertyCount); - generatedEmbeddings[vectorProperty] = [await floatTask.ConfigureAwait(false)]; - } - else - { - throw new InvalidOperationException( - $"The embedding generator configured on property '{vectorProperty.ModelName}' cannot produce an embedding of type '{typeof(Embedding).Name}' for the given input type."); - } - } - - using var connection = await this.GetConnectionAsync(cancellationToken).ConfigureAwait(false); - - await this.InternalUpsertBatchAsync(connection, [record], generatedEmbeddings, cancellationToken).ConfigureAwait(false); - } + public override Task UpsertAsync(TRecord record, CancellationToken cancellationToken = default) + => this.DoUpsertAsync([record], cancellationToken); /// - public override async Task UpsertAsync(IEnumerable records, CancellationToken cancellationToken = default) + public override Task UpsertAsync(IEnumerable records, CancellationToken cancellationToken = default) { Verify.NotNull(records); - IReadOnlyList? recordsList = null; - - // If an embedding generator is defined, invoke it once per property for all records. - Dictionary>>? generatedEmbeddings = null; - - var vectorPropertyCount = this._model.VectorProperties.Count; - for (var i = 0; i < vectorPropertyCount; i++) - { - var vectorProperty = this._model.VectorProperties[i]; - - if (SqliteModelBuilder.IsVectorPropertyTypeValidCore(vectorProperty.Type, out _)) - { - continue; - } - - // We have a vector property whose type isn't natively supported - we need to generate embeddings. - Debug.Assert(vectorProperty.EmbeddingGenerator is not null); - - // We have a property with embedding generation; materialize the records' enumerable if needed, to - // prevent multiple enumeration. - if (recordsList is null) - { - recordsList = records is IReadOnlyList r ? r : records.ToList(); - - if (recordsList.Count == 0) - { - return; - } - - records = recordsList; - } - - // TODO: Ideally we'd group together vector properties using the same generator (and with the same input and output properties), - // and generate embeddings for them in a single batch. That's some more complexity though. - if (vectorProperty.TryGenerateEmbeddings>(records, cancellationToken, out var floatTask)) - { - generatedEmbeddings ??= new Dictionary>>(vectorPropertyCount); - generatedEmbeddings[vectorProperty] = await floatTask.ConfigureAwait(false); - } - else - { - throw new InvalidOperationException( - $"The embedding generator configured on property '{vectorProperty.ModelName}' cannot produce an embedding of type '{typeof(Embedding).Name}' for the given input type."); - } - } - - using var connection = await this.GetConnectionAsync(cancellationToken).ConfigureAwait(false); - - await this.InternalUpsertBatchAsync(connection, records, generatedEmbeddings, cancellationToken).ConfigureAwait(false); + return this.DoUpsertAsync(records, cancellationToken); } /// @@ -631,27 +548,98 @@ private async IAsyncEnumerable InternalGetBatchAsync( } } - private async Task InternalUpsertBatchAsync( - SqliteConnection connection, - IEnumerable records, - Dictionary>>? generatedEmbeddings, - CancellationToken cancellationToken) + private async Task DoUpsertAsync(IEnumerable records, CancellationToken cancellationToken) { Verify.NotNull(records); - if (this._vectorPropertiesExist) + // With SQLite, we'll need to enumerate the records multiple times in almost all cases (e.g. because of the existence + // of two separate tables for data and vectors). To avoid multiple enumerations, we materialize the records into a list here. + var recordsList = records is IReadOnlyList r ? r : records.ToList(); + if (recordsList.Count == 0) { - // We're going to have to traverse the records multiple times, so materialize the enumerable if needed. - var recordsList = records is IReadOnlyList r ? r : records.ToList(); + return; + } + records = recordsList; + + // If an embedding generator is defined, invoke it once per property for all records. + Dictionary>>? generatedEmbeddings = null; + + var vectorPropertyCount = this._model.VectorProperties.Count; + for (var i = 0; i < vectorPropertyCount; i++) + { + var vectorProperty = this._model.VectorProperties[i]; - if (recordsList.Count == 0) + if (SqliteModelBuilder.IsVectorPropertyTypeValidCore(vectorProperty.Type, out _)) { - return; + continue; } - records = recordsList; + // We have a vector property whose type isn't natively supported - we need to generate embeddings. + Debug.Assert(vectorProperty.EmbeddingGenerator is not null); + + // TODO: Ideally we'd group together vector properties using the same generator (and with the same input and output properties), + // and generate embeddings for them in a single batch. That's some more complexity though. + if (vectorProperty.TryGenerateEmbeddings>(records, cancellationToken, out var floatTask)) + { + generatedEmbeddings ??= new Dictionary>>(vectorPropertyCount); + generatedEmbeddings[vectorProperty] = await floatTask.ConfigureAwait(false); + } + else + { + throw new InvalidOperationException( + $"The embedding generator configured on property '{vectorProperty.ModelName}' cannot produce an embedding of type '{typeof(Embedding).Name}' for the given input type."); + } + } + + var keyProperty = this._model.KeyProperty; + + using var connection = await this.GetConnectionAsync(cancellationToken).ConfigureAwait(false); + + using var dataCommand = SqliteCommandBuilder.BuildInsertCommand( + connection, + this._dataTableName, + this._model, + recordsList, + generatedEmbeddings, + data: true, + replaceIfExists: true); + + using (var reader = await connection.ExecuteWithErrorHandlingAsync( + this._collectionMetadata, + "updateData", + () => dataCommand.ExecuteReaderAsync(cancellationToken), + cancellationToken).ConfigureAwait(false)) + { + // If the key property is auto-generated, we need to read the generated keys from the database and inject them into the records + // (except for GUIDs which are generated client-side and have already been injected). + if (keyProperty is KeyPropertyModel { IsAutoGenerated: true } && keyProperty.Type != typeof(Guid)) + { + int? keyOrdinal = null; - var keyProperty = this._model.KeyProperty; + foreach (var record in recordsList) + { + switch (keyProperty.Type) + { + case var t when t == typeof(int) && keyProperty.GetValue(record) == 0: + keyOrdinal ??= reader.GetOrdinal(keyProperty.StorageName); + await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + keyProperty.SetValue(record, reader.GetFieldValue(keyOrdinal.Value)); + await reader.NextResultAsync(cancellationToken).ConfigureAwait(false); + continue; + case var t when t == typeof(long) && keyProperty.GetValue(record) == 0L: + keyOrdinal ??= reader.GetOrdinal(keyProperty.StorageName); + await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + keyProperty.SetValue(record, reader.GetFieldValue(keyOrdinal.Value)); + await reader.NextResultAsync(cancellationToken).ConfigureAwait(false); + continue; + } + } + } + } + + // We've inserted the main data records, now insert the records into the vector virtual table as well. + if (this._vectorPropertiesExist) + { var keys = recordsList.Select(r => keyProperty.GetValueAsObject(r)!).ToList(); // Deleting vector records first since current version of vector search extension @@ -670,9 +658,8 @@ await connection.ExecuteWithErrorHandlingAsync( using var vectorInsertCommand = SqliteCommandBuilder.BuildInsertCommand( connection, this._vectorTableName, - this._keyStorageName, this._model, - records, + recordsList, generatedEmbeddings, data: false); @@ -682,31 +669,6 @@ await connection.ExecuteWithErrorHandlingAsync( () => vectorInsertCommand.ExecuteNonQueryAsync(cancellationToken), cancellationToken).ConfigureAwait(false); } - - using var dataCommand = SqliteCommandBuilder.BuildInsertCommand( - connection, - this._dataTableName, - this._keyStorageName, - this._model, - records, - generatedEmbeddings, - data: true, - replaceIfExists: true); - - using var reader = await connection.ExecuteWithErrorHandlingAsync( - this._collectionMetadata, - "updateData", - () => dataCommand.ExecuteReaderAsync(cancellationToken), - cancellationToken).ConfigureAwait(false); - - while (await reader.ReadAsync(cancellationToken).ConfigureAwait(false)) - { - var key = reader.GetFieldValue(0); - - // TODO: Inject the generated keys into the record for autogenerated keys. - - await reader.NextResultAsync(cancellationToken).ConfigureAwait(false); - } } private Task InternalDeleteBatchAsync(SqliteConnection connection, SqliteWhereCondition condition, CancellationToken cancellationToken) diff --git a/dotnet/src/VectorData/SqliteVec/SqliteCommandBuilder.cs b/dotnet/src/VectorData/SqliteVec/SqliteCommandBuilder.cs index 526fa96919d6..7cc94a90da64 100644 --- a/dotnet/src/VectorData/SqliteVec/SqliteCommandBuilder.cs +++ b/dotnet/src/VectorData/SqliteVec/SqliteCommandBuilder.cs @@ -114,9 +114,8 @@ public static DbCommand BuildDropTableCommand(SqliteConnection connection, strin public static DbCommand BuildInsertCommand( SqliteConnection connection, string tableName, - string rowIdentifier, CollectionModel model, - IEnumerable records, + IReadOnlyList records, Dictionary>>? generatedEmbeddings, bool data, bool replaceIfExists = false) @@ -126,25 +125,34 @@ public static DbCommand BuildInsertCommand( var recordIndex = 0; - var properties = model.KeyProperties.Concat(data ? model.DataProperties : (IEnumerable)model.VectorProperties); + var properties = model.KeyProperties.Concat(data ? model.DataProperties : (IEnumerable)model.VectorProperties).ToList(); + var keyProperty = model.KeyProperty; + var isKeyPossiblyDatabaseGenerated = keyProperty.IsAutoGenerated && (keyProperty.Type == typeof(int) || keyProperty.Type == typeof(long)); foreach (var record in records) { - var rowIdentifierParameterName = GetParameterName(rowIdentifier, recordIndex); + var isRecordKeyDatabaseGenerated = isKeyPossiblyDatabaseGenerated + ? (keyProperty.Type == typeof(int) && keyProperty.GetValue(record) is var i && i == 0) + || (keyProperty.Type == typeof(long) && keyProperty.GetValue(record) is var l && l == 0L) + : false; sql.Append("INSERT"); - if (replaceIfExists) + if (replaceIfExists && !isRecordKeyDatabaseGenerated) { sql.Append(" OR REPLACE"); } sql.Append(" INTO ").AppendIdentifier(tableName).Append(" ("); -#pragma warning disable CA1851 // Possible multiple enumerations of 'IEnumerable' collection var propertyIndex = 0; foreach (var property in properties) { + if (property is KeyPropertyModel && isRecordKeyDatabaseGenerated) + { + continue; + } + if (propertyIndex++ > 0) { sql.Append(", "); @@ -160,45 +168,84 @@ public static DbCommand BuildInsertCommand( propertyIndex = 0; foreach (var property in properties) { - var parameterName = GetParameterName(property.StorageName, recordIndex); - - if (propertyIndex++ > 0) - { - sql.Append(", "); - } - - sql.Append(parameterName); - var value = property.GetValueAsObject(record); - if (property is VectorPropertyModel vectorProperty) + switch (property) { - if (generatedEmbeddings?[vectorProperty] is IReadOnlyList ge) + case KeyPropertyModel { IsAutoGenerated: true }: { - value = ((Embedding)ge[recordIndex]).Vector; + switch (value) + { + // Database ROWID generation. We don't specify the value in INSERT, and read it back later via RETURNING. + case int or long when isRecordKeyDatabaseGenerated: + continue; + + case Guid g when g == Guid.Empty: + // As SQLite has no built-in GUID generation, we generate client-side + // If the key is a Guid and auto-generated, generate a new Guid for any record where the key is empty. + if (keyProperty.IsAutoGenerated && keyProperty.Type == typeof(Guid)) + { +#if NET9_0_OR_GREATER + g = Guid.CreateVersion7(); +#else + g = Guid.NewGuid(); +#endif + value = g; + keyProperty.SetValue(record, g); + } + break; + + // We're configured for auto-generation but the user specified an explicit value. + case int or long or Guid: + break; + + default: + throw new UnreachableException(); + } + + break; } - value = value switch + case VectorPropertyModel vectorProperty: { - ReadOnlyMemory m => SqlitePropertyMapping.MapVectorForStorageModel(m), - Embedding e => SqlitePropertyMapping.MapVectorForStorageModel(e.Vector), - float[] a => SqlitePropertyMapping.MapVectorForStorageModel(a), - null => null, + if (generatedEmbeddings?[vectorProperty] is IReadOnlyList ge) + { + value = ((Embedding)ge[recordIndex]).Vector; + } + + value = value switch + { + ReadOnlyMemory m => SqlitePropertyMapping.MapVectorForStorageModel(m), + Embedding e => SqlitePropertyMapping.MapVectorForStorageModel(e.Vector), + float[] a => SqlitePropertyMapping.MapVectorForStorageModel(a), + null => null, + + _ => throw new InvalidOperationException($"Retrieved value for vector property '{property.StorageName}' which is not a ReadOnlyMemory ('{value?.GetType().Name}').") + }; + break; + } + } + + var parameterName = GetParameterName(property.StorageName, recordIndex); - _ => throw new InvalidOperationException($"Retrieved value for vector property '{property.StorageName}' which is not a ReadOnlyMemory ('{value?.GetType().Name}').") - }; + if (propertyIndex++ > 0) + { + sql.Append(", "); } + sql.Append(parameterName); + command.Parameters.Add(new SqliteParameter(parameterName, value ?? DBNull.Value)); } -#pragma warning restore CA1851 sql.AppendLine(")"); - sql - .Append("RETURNING ") - .AppendLine(rowIdentifier) - .AppendLine(";"); + if (isRecordKeyDatabaseGenerated) + { + sql.Append("RETURNING \"").Append(keyProperty.StorageName).Append('"'); + } + + sql.AppendLine(";"); recordIndex++; } diff --git a/dotnet/src/VectorData/SqliteVec/SqliteModelBuilder.cs b/dotnet/src/VectorData/SqliteVec/SqliteModelBuilder.cs index d771bc016f9b..3f2935fe68ab 100644 --- a/dotnet/src/VectorData/SqliteVec/SqliteModelBuilder.cs +++ b/dotnet/src/VectorData/SqliteVec/SqliteModelBuilder.cs @@ -18,11 +18,17 @@ internal class SqliteModelBuilder() : CollectionModelBuilder(s_modelBuildingOpti SupportsMultipleVectors = true, }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override bool SupportsKeyAutoGeneration(Type keyPropertyType) + => keyPropertyType == typeof(Guid) || keyPropertyType == typeof(int) || keyPropertyType == typeof(long); + + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "int, long, string, Guid"; + var type = keyProperty.Type; - return type == typeof(int) || type == typeof(long) || type == typeof(string) || type == typeof(Guid); + if (type != typeof(int) && type != typeof(long) && type != typeof(string) && type != typeof(Guid)) + { + throw new NotSupportedException($"The property type '{type.FullName}' is not supported for key properties by the SqliteVec provider. Supported types are: int, long, string, Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/CollectionModelBuilder.cs b/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/CollectionModelBuilder.cs index 1ffc9f52b4a3..79af64cae8c5 100644 --- a/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/CollectionModelBuilder.cs +++ b/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/CollectionModelBuilder.cs @@ -152,6 +152,7 @@ protected virtual void ProcessTypeProperties(Type type, VectorStoreCollectionDef if (clrProperty.GetCustomAttribute() is { } keyAttribute) { var keyProperty = new KeyPropertyModel(clrProperty.Name, clrProperty.PropertyType); + keyProperty.IsAutoGenerated = keyAttribute.IsAutoGenerated ?? this.SupportsKeyAutoGeneration(keyProperty.Type); this.KeyProperties.Add(keyProperty); storageName = keyAttribute.StorageName; property = keyProperty; @@ -264,7 +265,7 @@ protected virtual void ProcessRecordDefinition(VectorStoreCollectionDefinition d switch (definitionProperty) { case VectorStoreKeyProperty definitionKeyProperty: - var keyProperty = new KeyPropertyModel(definitionKeyProperty.Name, propertyType!); + var keyProperty = new KeyPropertyModel(definitionKeyProperty.Name, propertyType); this.KeyProperties.Add(keyProperty); this.PropertyMap.Add(definitionKeyProperty.Name, keyProperty); property = keyProperty; @@ -297,6 +298,8 @@ protected virtual void ProcessRecordDefinition(VectorStoreCollectionDefinition d $"Property '{property.ModelName}' is present in the {nameof(VectorStoreCollectionDefinition)} as a key property, but the .NET property on type '{type?.Name}' has an incompatible attribute."); } + keyPropertyModel.IsAutoGenerated = definitionKeyProperty.IsAutoGenerated ?? this.SupportsKeyAutoGeneration(keyPropertyModel.Type); + break; case VectorStoreDataProperty definitionDataProperty: @@ -463,15 +466,17 @@ protected virtual void ValidateProperty(PropertyModel propertyModel, VectorStore switch (propertyModel) { case KeyPropertyModel keyProperty: - if (!this.IsKeyPropertyTypeValid(keyProperty.Type, out var supportedTypes)) + if (keyProperty.IsAutoGenerated && !this.SupportsKeyAutoGeneration(keyProperty.Type)) { throw new NotSupportedException( - $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: {supportedTypes}."); + $"Property '{keyProperty.ModelName}' is configured for auto-generation, but key properties of type '{keyProperty.Type.Name}' do not support auto-generation."); } + + this.ValidateKeyProperty(keyProperty); break; case DataPropertyModel dataProperty: - if (!this.IsDataPropertyTypeValid(dataProperty.Type, out supportedTypes)) + if (!this.IsDataPropertyTypeValid(dataProperty.Type, out var supportedTypes)) { throw new NotSupportedException( $"Property '{dataProperty.ModelName}' has unsupported type '{type.Name}'. Data properties must be one of the supported types: {supportedTypes}."); @@ -522,10 +527,19 @@ protected virtual void ValidateProperty(PropertyModel propertyModel, VectorStore } } + /// + /// Configures auto-generation for the given key property. + /// Defaults to configuring key properties as auto-generated, and throwing if auto-generation is requested for + /// any other type. + /// + protected virtual bool SupportsKeyAutoGeneration(Type keyPropertyType) + => keyPropertyType == typeof(Guid); + /// /// Validates that the .NET type for a key property is supported by the provider. /// - protected abstract bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes); + /// true if the provider supports auto-generating keys of the specified type; otherwise, false. + protected abstract void ValidateKeyProperty(KeyPropertyModel keyProperty); /// /// Validates that the .NET type for a data property is supported by the provider. diff --git a/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/KeyPropertyModel.cs b/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/KeyPropertyModel.cs index 341439020ce9..76f90822f7a5 100644 --- a/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/KeyPropertyModel.cs +++ b/dotnet/src/VectorData/VectorData.Abstractions/ProviderServices/KeyPropertyModel.cs @@ -12,7 +12,12 @@ namespace Microsoft.Extensions.VectorData.ProviderServices; [Experimental("MEVD9001")] public class KeyPropertyModel(string modelName, Type type) : PropertyModel(modelName, type) { + /// + /// Gets or sets whether this key property's value is auto-generated or not. + /// + public bool IsAutoGenerated { get; set; } + /// public override string ToString() - => $"{this.ModelName} (Key, {this.Type.Name})"; + => $"{this.ModelName} (Key, {this.Type.Name}{(this.IsAutoGenerated ? ", auto-generated" : "")})"; } diff --git a/dotnet/src/VectorData/VectorData.Abstractions/RecordAttributes/VectorStoreKeyAttribute.cs b/dotnet/src/VectorData/VectorData.Abstractions/RecordAttributes/VectorStoreKeyAttribute.cs index 488bab5511d8..26de020ba26d 100644 --- a/dotnet/src/VectorData/VectorData.Abstractions/RecordAttributes/VectorStoreKeyAttribute.cs +++ b/dotnet/src/VectorData/VectorData.Abstractions/RecordAttributes/VectorStoreKeyAttribute.cs @@ -20,4 +20,12 @@ public sealed class VectorStoreKeyAttribute : Attribute /// For example, the property name might be "MyProperty" and the storage name might be "my_property". /// public string? StorageName { get; init; } + + /// + /// Gets or sets whether this key property's value is auto-generated or not. + /// + /// + /// The availability of auto-generated properties - as well as the .NET types supported for them - varies across provider implementations. + /// + public bool? IsAutoGenerated { get; set; } } diff --git a/dotnet/src/VectorData/VectorData.Abstractions/RecordDefinition/VectorStoreKeyProperty.cs b/dotnet/src/VectorData/VectorData.Abstractions/RecordDefinition/VectorStoreKeyProperty.cs index 12c849d16c9c..2c2a6f6262f0 100644 --- a/dotnet/src/VectorData/VectorData.Abstractions/RecordDefinition/VectorStoreKeyProperty.cs +++ b/dotnet/src/VectorData/VectorData.Abstractions/RecordDefinition/VectorStoreKeyProperty.cs @@ -21,4 +21,12 @@ public VectorStoreKeyProperty(string name, Type? type = null) : base(name, type) { } + + /// + /// Gets or sets whether this key property's value is auto-generated or not. + /// + /// + /// The availability of auto-generated properties - as well as the .NET types supported for them - varies across provider implementations. + /// + public bool? IsAutoGenerated { get; set; } } diff --git a/dotnet/src/VectorData/Weaviate/WeaviateCollection.cs b/dotnet/src/VectorData/Weaviate/WeaviateCollection.cs index ac7005f43e2d..de1d482633b5 100644 --- a/dotnet/src/VectorData/Weaviate/WeaviateCollection.cs +++ b/dotnet/src/VectorData/Weaviate/WeaviateCollection.cs @@ -304,7 +304,18 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio } } - var jsonObjects = records.Select((record, i) => this._mapper.MapFromDataToStorageModel(record, i, generatedEmbeddings)).ToList(); + var keyProperty = this._model.KeyProperty; + var jsonObjects = new List(); + var recordIndex = 0; + foreach (var record in records) + { + if (keyProperty.IsAutoGenerated && keyProperty.GetValue(record) == Guid.Empty) + { + keyProperty.SetValue(record, Guid.NewGuid()); + } + + jsonObjects.Add(this._mapper.MapFromDataToStorageModel(record, recordIndex++, generatedEmbeddings)); + } if (jsonObjects.Count == 0) { diff --git a/dotnet/src/VectorData/Weaviate/WeaviateModelBuilder.cs b/dotnet/src/VectorData/Weaviate/WeaviateModelBuilder.cs index 5be7bfa1d5e5..9a1f429849de 100644 --- a/dotnet/src/VectorData/Weaviate/WeaviateModelBuilder.cs +++ b/dotnet/src/VectorData/Weaviate/WeaviateModelBuilder.cs @@ -23,11 +23,13 @@ private static CollectionModelBuildingOptions GetModelBuildingOptions(bool hasNa }; } - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "Guid"; - - return type == typeof(Guid); + if (keyProperty.Type != typeof(Guid)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{keyProperty.Type.Name}'. Key properties must be of type Guid."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) diff --git a/dotnet/test/VectorData/CosmosMongoDB.ConformanceTests/TypeTests/CosmosMongoKeyTypeTests.cs b/dotnet/test/VectorData/CosmosMongoDB.ConformanceTests/TypeTests/CosmosMongoKeyTypeTests.cs index 78c86e17ed8d..62533ade31e3 100644 --- a/dotnet/test/VectorData/CosmosMongoDB.ConformanceTests/TypeTests/CosmosMongoKeyTypeTests.cs +++ b/dotnet/test/VectorData/CosmosMongoDB.ConformanceTests/TypeTests/CosmosMongoKeyTypeTests.cs @@ -13,16 +13,16 @@ public class CosmosMongoKeyTypeTests(CosmosMongoKeyTypeTests.Fixture fixture) : KeyTypeTests(fixture), IClassFixture { [ConditionalFact] - public virtual Task ObjectId() => this.Test(new("652f8c3e8f9b2c1a4d3e6a7b"), new("b7a6e3d4a1c2b9f8e3c8f256")); + public virtual Task ObjectId() => this.Test(new("652f8c3e8f9b2c1a4d3e6a7b"), supportsAutoGeneration: true); [ConditionalFact] public virtual Task String() => this.Test("foo", "bar"); [ConditionalFact] - public virtual Task Int() => this.Test(8, 9); + public virtual Task Int() => this.Test(8); [ConditionalFact] - public virtual Task Long() => this.Test(8L, 9L); + public virtual Task Long() => this.Test(8L); public new class Fixture : KeyTypeTests.Fixture { diff --git a/dotnet/test/VectorData/InMemory.ConformanceTests/TypeTests/InMemoryKeyTypeTests.cs b/dotnet/test/VectorData/InMemory.ConformanceTests/TypeTests/InMemoryKeyTypeTests.cs index e7d94a95fece..d0d8f4e3bf8c 100644 --- a/dotnet/test/VectorData/InMemory.ConformanceTests/TypeTests/InMemoryKeyTypeTests.cs +++ b/dotnet/test/VectorData/InMemory.ConformanceTests/TypeTests/InMemoryKeyTypeTests.cs @@ -22,13 +22,13 @@ public class InMemoryKeyTypeTests(InMemoryKeyTypeTests.Fixture fixture) [ConditionalFact] public virtual Task String() => this.Test("foo", "bar"); - protected override async Task Test(TKey key1, TKey key2) + protected override async Task Test(TKey key1, TKey key2, bool supportsAutoGeneration = false) { - await base.Test(key1, key2); + await base.Test(key1, key2, supportsAutoGeneration); // For InMemory, delete the collection, otherwise the next test that runs will fail because the collection // already exists but with the previous key type. - using var collection = fixture.CreateCollection(); + using var collection = fixture.CreateCollection(supportsAutoGeneration); await collection.EnsureCollectionDeletedAsync(); } diff --git a/dotnet/test/VectorData/MongoDB.ConformanceTests/TypeTests/MongoKeyTypeTests.cs b/dotnet/test/VectorData/MongoDB.ConformanceTests/TypeTests/MongoKeyTypeTests.cs index 19d847e1c0b2..89989a9e128c 100644 --- a/dotnet/test/VectorData/MongoDB.ConformanceTests/TypeTests/MongoKeyTypeTests.cs +++ b/dotnet/test/VectorData/MongoDB.ConformanceTests/TypeTests/MongoKeyTypeTests.cs @@ -13,16 +13,16 @@ public class MongoKeyTypeTests(MongoKeyTypeTests.Fixture fixture) : KeyTypeTests(fixture), IClassFixture { [ConditionalFact] - public virtual Task ObjectId() => this.Test(new("652f8c3e8f9b2c1a4d3e6a7b"), new("b7a6e3d4a1c2b9f8e3c8f256")); + public virtual Task ObjectId() => this.Test(new("652f8c3e8f9b2c1a4d3e6a7b"), supportsAutoGeneration: true); [ConditionalFact] public virtual Task String() => this.Test("foo", "bar"); [ConditionalFact] - public virtual Task Int() => this.Test(8, 9); + public virtual Task Int() => this.Test(8); [ConditionalFact] - public virtual Task Long() => this.Test(8L, 9L); + public virtual Task Long() => this.Test(8L); public new class Fixture : KeyTypeTests.Fixture { diff --git a/dotnet/test/VectorData/PgVector.ConformanceTests/Support/PostgresTestStore.cs b/dotnet/test/VectorData/PgVector.ConformanceTests/Support/PostgresTestStore.cs index 9deae8cf8b73..c9e19cbdff43 100644 --- a/dotnet/test/VectorData/PgVector.ConformanceTests/Support/PostgresTestStore.cs +++ b/dotnet/test/VectorData/PgVector.ConformanceTests/Support/PostgresTestStore.cs @@ -14,7 +14,7 @@ internal sealed class PostgresTestStore : TestStore public static PostgresTestStore Instance { get; } = new(); private static readonly PostgreSqlContainer s_container = new PostgreSqlBuilder() - .WithImage("pgvector/pgvector:pg17") + .WithImage("pgvector/pgvector:pg18") .Build(); private NpgsqlDataSource? _dataSource; diff --git a/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresKeyTypeTests.cs b/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresKeyTypeTests.cs index b62a138ed0a2..659b4d8bdcb6 100644 --- a/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresKeyTypeTests.cs +++ b/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresKeyTypeTests.cs @@ -12,10 +12,10 @@ public class PostgresKeyTypeTests(PostgresKeyTypeTests.Fixture fixture) : KeyTypeTests(fixture), IClassFixture { [ConditionalFact] - public virtual Task Int() => this.Test(8, 9); + public virtual Task Int() => this.Test(8, supportsAutoGeneration: true); [ConditionalFact] - public virtual Task Long() => this.Test(8L, 9L); + public virtual Task Long() => this.Test(8L, supportsAutoGeneration: true); [ConditionalFact] public virtual Task String() => this.Test("foo", "bar"); diff --git a/dotnet/test/VectorData/PgVector.UnitTests/PostgresSqlBuilderTests.cs b/dotnet/test/VectorData/PgVector.UnitTests/PostgresSqlBuilderTests.cs index 810e1257b898..43ae2533173c 100644 --- a/dotnet/test/VectorData/PgVector.UnitTests/PostgresSqlBuilderTests.cs +++ b/dotnet/test/VectorData/PgVector.UnitTests/PostgresSqlBuilderTests.cs @@ -51,7 +51,7 @@ public void TestBuildCreateTableCommand(bool ifNotExists) var model = new PostgresModelBuilder().BuildDynamic(recordDefinition, defaultEmbeddingGenerator: null); - var sql = PostgresSqlBuilder.BuildCreateTableSql("public", "testcollection", model, ifNotExists: ifNotExists); + var sql = PostgresSqlBuilder.BuildCreateTableSql("public", "testcollection", model, pgVersion: new Version(18, 0), ifNotExists: ifNotExists); // Check for expected properties; integration tests will validate the actual SQL. Assert.Contains("\"public\".\"testcollection\" (", sql); @@ -198,10 +198,12 @@ public void TestBuildUpsertCommand() ["embedding1"] = new ReadOnlyMemory(s_vector), }; - using var command = new NpgsqlCommand(); - var cmdInfo = PostgresSqlBuilder.BuildUpsertCommand(command, "public", "testcollection", model, [record], generatedEmbeddings: null); + using var batch = new NpgsqlBatch(); + _ = PostgresSqlBuilder.BuildUpsertCommand(batch, "public", "testcollection", model, [record], generatedEmbeddings: null); // Check for expected properties; integration tests will validate the actual SQL. + Assert.Single(batch.BatchCommands); + var command = batch.BatchCommands[0]; Assert.Contains("INSERT INTO \"public\".\"testcollection\" (", command.CommandText); Assert.Contains("ON CONFLICT (\"id\")", command.CommandText); Assert.Contains("DO UPDATE SET", command.CommandText); diff --git a/dotnet/test/VectorData/Qdrant.ConformanceTests/TypeTests/QdrantKeyTypeTests.cs b/dotnet/test/VectorData/Qdrant.ConformanceTests/TypeTests/QdrantKeyTypeTests.cs index 2334f41a81a2..6594bddb3e35 100644 --- a/dotnet/test/VectorData/Qdrant.ConformanceTests/TypeTests/QdrantKeyTypeTests.cs +++ b/dotnet/test/VectorData/Qdrant.ConformanceTests/TypeTests/QdrantKeyTypeTests.cs @@ -12,7 +12,7 @@ public class QdrantKeyTypeTests(QdrantKeyTypeTests.Fixture fixture) : KeyTypeTests(fixture), IClassFixture { [ConditionalFact] - public virtual Task ULong() => this.Test(8UL, 9UL); + public virtual Task ULong() => this.Test(8UL); public new class Fixture : KeyTypeTests.Fixture { diff --git a/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisHashSetKeyTypeTests.cs b/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisHashSetKeyTypeTests.cs index bc84e34f7fed..bbf40144fea9 100644 --- a/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisHashSetKeyTypeTests.cs +++ b/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisHashSetKeyTypeTests.cs @@ -23,10 +23,10 @@ public class RedisHashSetKeyTypeTests(RedisHashSetKeyTypeTests.Fixture fixture) // Redis doesn't seem to reliably delete the collection: when running multiple tests that delete and recreate the collection with different key types, // we seem to get key values from the previous collection despite having deleted and recreated it. So we uniquify the collection name instead. - public override VectorStoreCollection> CreateCollection() - => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition()); + public override VectorStoreCollection> CreateCollection(bool? withAutoGeneration) + => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition(withAutoGeneration)); - public override VectorStoreCollection> CreateDynamicCollection() - => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition()); + public override VectorStoreCollection> CreateDynamicCollection(bool withAutoGeneration) + => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition(withAutoGeneration)); } } diff --git a/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisJsonKeyTypeTests.cs b/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisJsonKeyTypeTests.cs index 2bf539fa9124..ecc2030a65e2 100644 --- a/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisJsonKeyTypeTests.cs +++ b/dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisJsonKeyTypeTests.cs @@ -23,10 +23,10 @@ public class RedisJsonKeyTypeTests(RedisJsonKeyTypeTests.Fixture fixture) // Redis doesn't seem to reliably delete the collection: when running multiple tests that delete and recreate the collection with different key types, // we seem to get key values from the previous collection despite having deleted and recreated it. So we uniquify the collection name instead. - public override VectorStoreCollection> CreateCollection() - => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition()); + public override VectorStoreCollection> CreateCollection(bool? withAutoGeneration) + => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition(withAutoGeneration)); - public override VectorStoreCollection> CreateDynamicCollection() - => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition()); + public override VectorStoreCollection> CreateDynamicCollection(bool withAutoGeneration) + => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition(withAutoGeneration)); } } diff --git a/dotnet/test/VectorData/SqlServer.ConformanceTests/SqlServerCommandBuilderTests.cs b/dotnet/test/VectorData/SqlServer.ConformanceTests/SqlServerCommandBuilderTests.cs index a8100b15c75c..126d681aa2e8 100644 --- a/dotnet/test/VectorData/SqlServer.ConformanceTests/SqlServerCommandBuilderTests.cs +++ b/dotnet/test/VectorData/SqlServer.ConformanceTests/SqlServerCommandBuilderTests.cs @@ -125,7 +125,7 @@ public void CreateTable(bool ifNotExists) """ BEGIN CREATE TABLE [schema].[table] ( - [id] BIGINT NOT NULL, + [id] BIGINT IDENTITY, [simpleName] NVARCHAR(MAX), [with space] INT, [embedding] VECTOR(10), @@ -143,59 +143,11 @@ PRIMARY KEY ([id]) } [Fact] - public void MergeIntoSingle() + public void Upsert() { var model = BuildModel( [ - new VectorStoreKeyProperty("id", typeof(long)), - new VectorStoreDataProperty("simpleString", typeof(string)), - new VectorStoreDataProperty("simpleInt", typeof(int)), - new VectorStoreVectorProperty("embedding", typeof(ReadOnlyMemory), 10) - ]); - - using SqlConnection connection = CreateConnection(); - using SqlCommand command = SqlServerCommandBuilder.MergeIntoSingle( - connection, "schema", "table", model, - new Dictionary - { - { "id", null }, - { "simpleString", "nameValue" }, - { "simpleInt", 134 }, - { "embedding", new ReadOnlyMemory([10.0f]) } - }, - generatedEmbeddings: null); - - string expectedCommand = - """" - MERGE INTO [schema].[table] AS t - USING (VALUES (@id_0,@simpleString_1,@simpleInt_2,@embedding_3)) AS s ([id],[simpleString],[simpleInt],[embedding]) - ON (t.[id] = s.[id]) - WHEN MATCHED THEN - UPDATE SET t.[simpleString] = s.[simpleString],t.[simpleInt] = s.[simpleInt],t.[embedding] = s.[embedding] - WHEN NOT MATCHED THEN - INSERT ([id],[simpleString],[simpleInt],[embedding]) - VALUES (s.[id],s.[simpleString],s.[simpleInt],s.[embedding]) - OUTPUT inserted.[id]; - """"; - - Assert.Equal(expectedCommand, command.CommandText, ignoreLineEndingDifferences: true); - Assert.Equal("@id_0", command.Parameters[0].ParameterName); - Assert.Equal(DBNull.Value, command.Parameters[0].Value); - Assert.Equal("@simpleString_1", command.Parameters[1].ParameterName); - Assert.Equal("nameValue", command.Parameters[1].Value); - Assert.Equal("@simpleInt_2", command.Parameters[2].ParameterName); - Assert.Equal(134, command.Parameters[2].Value); - Assert.Equal("@embedding_3", command.Parameters[3].ParameterName); - var vector = Assert.IsType>(command.Parameters[3].Value); - Assert.Equal([10], vector.Memory.ToArray()); - } - - [Fact] - public void MergeIntoMany() - { - var model = BuildModel( - [ - new VectorStoreKeyProperty("id", typeof(long)), + new VectorStoreKeyProperty("id", typeof(long)) { IsAutoGenerated = false }, new VectorStoreDataProperty("simpleString", typeof(string)), new VectorStoreDataProperty("simpleInt", typeof(int)), new VectorStoreVectorProperty("embedding", typeof(ReadOnlyMemory), 10) @@ -222,23 +174,31 @@ public void MergeIntoMany() using SqlConnection connection = CreateConnection(); using SqlCommand command = connection.CreateCommand(); - Assert.True(SqlServerCommandBuilder.MergeIntoMany(command, "schema", "table", model, records, firstRecordIndex: 0, generatedEmbeddings: null)); + Assert.True(SqlServerCommandBuilder.Upsert(command, "schema", "table", model, records, firstRecordIndex: 0, generatedEmbeddings: null)); string expectedCommand = """" - DECLARE @InsertedKeys TABLE (KeyColumn BIGINT); MERGE INTO [schema].[table] AS t - USING (VALUES - (@id_0,@simpleString_1,@simpleInt_2,@embedding_3), - (@id_4,@simpleString_5,@simpleInt_6,@embedding_7)) AS s ([id],[simpleString],[simpleInt],[embedding]) + USING (VALUES (@id_0,@simpleString_1,@simpleInt_2,@embedding_3)) AS s ([id],[simpleString],[simpleInt],[embedding]) + ON (t.[id] = s.[id]) + WHEN MATCHED THEN + UPDATE SET t.[simpleString] = s.[simpleString],t.[simpleInt] = s.[simpleInt],t.[embedding] = s.[embedding] + WHEN NOT MATCHED THEN + INSERT ([id],[simpleString],[simpleInt],[embedding]) + VALUES (s.[id],s.[simpleString],s.[simpleInt],s.[embedding]) + OUTPUT inserted.[id]; + + MERGE INTO [schema].[table] AS t + USING (VALUES (@id_4,@simpleString_5,@simpleInt_6,@embedding_7)) AS s ([id],[simpleString],[simpleInt],[embedding]) ON (t.[id] = s.[id]) WHEN MATCHED THEN UPDATE SET t.[simpleString] = s.[simpleString],t.[simpleInt] = s.[simpleInt],t.[embedding] = s.[embedding] WHEN NOT MATCHED THEN INSERT ([id],[simpleString],[simpleInt],[embedding]) VALUES (s.[id],s.[simpleString],s.[simpleInt],s.[embedding]) - OUTPUT inserted.[id] INTO @InsertedKeys (KeyColumn); - SELECT KeyColumn FROM @InsertedKeys; + OUTPUT inserted.[id]; + + """"; Assert.Equal(expectedCommand, command.CommandText, ignoreLineEndingDifferences: true); diff --git a/dotnet/test/VectorData/SqlServer.ConformanceTests/TypeTests/SqlServerKeyTypeTests.cs b/dotnet/test/VectorData/SqlServer.ConformanceTests/TypeTests/SqlServerKeyTypeTests.cs index de130da41cfc..27d5d2019eb2 100644 --- a/dotnet/test/VectorData/SqlServer.ConformanceTests/TypeTests/SqlServerKeyTypeTests.cs +++ b/dotnet/test/VectorData/SqlServer.ConformanceTests/TypeTests/SqlServerKeyTypeTests.cs @@ -12,10 +12,10 @@ public class SqlServerKeyTypeTests(SqlServerKeyTypeTests.Fixture fixture) : KeyTypeTests(fixture), IClassFixture { [ConditionalFact] - public virtual Task Int() => this.Test(8, 9); + public virtual Task Int() => this.Test(8, supportsAutoGeneration: true); [ConditionalFact] - public virtual Task Long() => this.Test(8L, 9L); + public virtual Task Long() => this.Test(8L, supportsAutoGeneration: true); [ConditionalFact] public virtual Task String() => this.Test("foo", "bar"); diff --git a/dotnet/test/VectorData/SqliteVec.ConformanceTests/TypeTests/SqliteKeyTypeTests.cs b/dotnet/test/VectorData/SqliteVec.ConformanceTests/TypeTests/SqliteKeyTypeTests.cs index a12524f40100..8561d6eaba9e 100644 --- a/dotnet/test/VectorData/SqliteVec.ConformanceTests/TypeTests/SqliteKeyTypeTests.cs +++ b/dotnet/test/VectorData/SqliteVec.ConformanceTests/TypeTests/SqliteKeyTypeTests.cs @@ -12,10 +12,10 @@ public class SqliteKeyTypeTests(SqliteKeyTypeTests.Fixture fixture) : KeyTypeTests(fixture), IClassFixture { [ConditionalFact] - public virtual Task Int() => this.Test(8, 9); + public virtual Task Int() => this.Test(8, supportsAutoGeneration: true); [ConditionalFact] - public virtual Task Long() => this.Test(8L, 9L); + public virtual Task Long() => this.Test(8L, supportsAutoGeneration: true); [ConditionalFact] public virtual Task String() => this.Test("foo", "bar"); diff --git a/dotnet/test/VectorData/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs b/dotnet/test/VectorData/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs index bfbda2993930..eea00cbbaa10 100644 --- a/dotnet/test/VectorData/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs +++ b/dotnet/test/VectorData/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs @@ -110,15 +110,14 @@ public void ItBuildsDropTableCommand() [Theory] [InlineData(true)] [InlineData(false)] - public void ItBuildsInsertCommand(bool replaceIfExists) + public void ItBuildsInsertCommand_without_autogenerated_key(bool replaceIfExists) { // Arrange const string TableName = "TestTable"; - const string RowIdentifier = "Id"; var model = BuildModel( [ - new VectorStoreKeyProperty("Id", typeof(string)), + new VectorStoreKeyProperty("Id", typeof(int)), new VectorStoreDataProperty("Name", typeof(string)), new VectorStoreDataProperty("Age", typeof(string)), new VectorStoreDataProperty("Address", typeof(string)), @@ -126,20 +125,19 @@ public void ItBuildsInsertCommand(bool replaceIfExists) var records = new List> { - new() { ["Id"] = "IdValue1", ["Name"] = "NameValue1", ["Age"] = "AgeValue1", ["Address"] = "AddressValue1" }, - new() { ["Id"] = "IdValue2", ["Name"] = "NameValue2", ["Age"] = "AgeValue2", ["Address"] = "AddressValue2" }, + new() { ["Id"] = 1, ["Name"] = "NameValue1", ["Age"] = "AgeValue1", ["Address"] = "AddressValue1" }, + new() { ["Id"] = 2, ["Name"] = "NameValue2", ["Age"] = "AgeValue2", ["Address"] = "AddressValue2" }, }; // Act var command = SqliteCommandBuilder.BuildInsertCommand( this._connection, TableName, - RowIdentifier, model, records, generatedEmbeddings: null, data: true, - replaceIfExists); + replaceIfExists: replaceIfExists); // Assert Assert.Equal(replaceIfExists, command.CommandText.Contains("OR REPLACE")); @@ -147,10 +145,10 @@ public void ItBuildsInsertCommand(bool replaceIfExists) Assert.Contains($"INTO \"{TableName}\" (\"Id\", \"Name\", \"Age\", \"Address\")", command.CommandText); Assert.Contains("VALUES (@Id0, @Name0, @Age0, @Address0)", command.CommandText); Assert.Contains("VALUES (@Id1, @Name1, @Age1, @Address1)", command.CommandText); - Assert.Contains("RETURNING Id", command.CommandText); + Assert.DoesNotContain("RETURNING", command.CommandText); Assert.Equal("@Id0", command.Parameters[0].ParameterName); - Assert.Equal("IdValue1", command.Parameters[0].Value); + Assert.Equal(1, command.Parameters[0].Value); Assert.Equal("@Name0", command.Parameters[1].ParameterName); Assert.Equal("NameValue1", command.Parameters[1].Value); @@ -162,7 +160,7 @@ public void ItBuildsInsertCommand(bool replaceIfExists) Assert.Equal("AddressValue1", command.Parameters[3].Value); Assert.Equal("@Id1", command.Parameters[4].ParameterName); - Assert.Equal("IdValue2", command.Parameters[4].Value); + Assert.Equal(2, command.Parameters[4].Value); Assert.Equal("@Name1", command.Parameters[5].ParameterName); Assert.Equal("NameValue2", command.Parameters[5].Value); @@ -174,6 +172,53 @@ public void ItBuildsInsertCommand(bool replaceIfExists) Assert.Equal("AddressValue2", command.Parameters[7].Value); } + [Theory] + [InlineData(true)] + [InlineData(false)] + public void ItBuildsInsertCommand_with_autogenerated_key(bool replaceIfExists) + { + // Arrange + const string TableName = "TestTable"; + + var model = BuildModel( + [ + new VectorStoreKeyProperty("Id", typeof(int)), + new VectorStoreDataProperty("Name", typeof(string)), + new VectorStoreDataProperty("Age", typeof(string)), + new VectorStoreDataProperty("Address", typeof(string)), + ]); + + var records = new List> + { + new() { ["Id"] = default(int), ["Name"] = "NameValue1", ["Age"] = "AgeValue1", ["Address"] = "AddressValue1" }, + new() { ["Id"] = default(int), ["Name"] = "NameValue2", ["Age"] = "AgeValue2", ["Address"] = "AddressValue2" }, + }; + + // Act + var command = SqliteCommandBuilder.BuildInsertCommand( + this._connection, + TableName, + model, + records, + generatedEmbeddings: null, + data: true, + replaceIfExists: replaceIfExists); + + // Assert + Assert.DoesNotContain("OR REPLACE", command.CommandText); + + Assert.Contains($"INTO \"{TableName}\" (\"Name\", \"Age\", \"Address\")", command.CommandText); + Assert.Contains("VALUES (@Name0, @Age0, @Address0)", command.CommandText); + Assert.Contains("VALUES (@Name1, @Age1, @Address1)", command.CommandText); + Assert.Contains("RETURNING \"Id\"", command.CommandText); + + Assert.Equal("@Name0", command.Parameters[0].ParameterName); + Assert.Equal("NameValue1", command.Parameters[0].Value); + + Assert.Equal("@Name1", command.Parameters[3].ParameterName); + Assert.Equal("NameValue2", command.Parameters[3].Value); + } + [Theory] [InlineData(null)] [InlineData("")] diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs index ec3ebf22b823..fb9db07dac62 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs @@ -9,51 +9,75 @@ namespace VectorData.ConformanceTests.TypeTests; public abstract class KeyTypeTests(KeyTypeTests.Fixture fixture) { - // All MEVD providers are expected to support Guid keys (possibly by storing them as strings). + // All MEVD providers are expected to support Guid keys (possibly by storing them as strings), including + // auto-generation. // This allows upper layers such as Microsoft.Extensions.DataIngestion to use Guid keys consistently. [ConditionalFact] - public virtual Task Guid() => this.Test(new Guid("603840bf-cf91-4521-8b8e-8b6a2e75910a"), new Guid("f507c6a2-43bd-4d8f-8656-890f2cdaf299")); - - protected virtual async Task Test(TKey key1, TKey key2) + public virtual Task Guid() + => this.Test( + new Guid("603840bf-cf91-4521-8b8e-8b6a2e75910a"), + supportsAutoGeneration: true); + + protected virtual Task Test(TKey key, bool supportsAutoGeneration = false) + where TKey : struct + => this.Test(key, default!, supportsAutoGeneration: supportsAutoGeneration); + + // Note that we do not currently support testing auto generation for reference types, since + // no such case currently exists in a known provider. As a result we require a second key + // value. + protected virtual Task Test(TKey key1, TKey key2) + where TKey : class + => this.Test(key1, key2, supportsAutoGeneration: false); + + protected virtual async Task Test( + TKey key1, + TKey key2, + bool supportsAutoGeneration = false) where TKey : notnull { Assert.NotEqual(key1, key2); - using var collection = fixture.CreateCollection(); - - await collection.EnsureCollectionDeletedAsync(); - await collection.EnsureCollectionExistsAsync(); - - var record = new Record - { - Key = key1, - Int = 8, - Vector = new ReadOnlyMemory([1, 2, 3]) - }; + using var collection = fixture.CreateCollection(withAutoGeneration: false); - var nextRecord = new Record { - Key = key2, - Int = 9, - Vector = new ReadOnlyMemory([3, 2, 1]) - }; + await collection.EnsureCollectionDeletedAsync(); + await collection.EnsureCollectionExistsAsync(); - await collection.UpsertAsync(record); - await collection.UpsertAsync([record, nextRecord]); - await fixture.TestStore.WaitForDataAsync(collection, recordCount: 2); - - var result = await collection.GetAsync(key1); + var record = new Record + { + Key = key1, + Int = 8, + Vector = new ReadOnlyMemory([1, 2, 3]) + }; - Assert.NotNull(result); - Assert.Equal(key1, result.Key); - Assert.Equal(8, result.Int); + var nextRecord = new Record + { + Key = key2, + Int = 9, + Vector = new ReadOnlyMemory([3, 2, 1]) + }; - var results = await collection.GetAsync([key1, key2]).ToListAsync(); - Assert.Equal(2, results.Count); - var firstRecord = Assert.Single(results, r => r.Key.Equals(key1)); - Assert.Equal(8, firstRecord.Int); - var secondRecord = Assert.Single(results, r => r.Key.Equals(key2)); - Assert.Equal(9, secondRecord.Int); + await collection.UpsertAsync(record); + // Exercise multi-record plus updating existing record + await collection.UpsertAsync([record, nextRecord]); + await fixture.TestStore.WaitForDataAsync(collection, recordCount: 2); + + // Single record get + var result = await collection.GetAsync(key1); + Assert.NotNull(result); + Assert.Equal(key1, result.Key); + Assert.Equal(8, result.Int); + + // Multiple record get + // Also ensures that the second record - with the default key value - got properly inserted and did not trigger auto-generation + // (as we haven't configured it). + var results = await collection.GetAsync([key1, key2]).ToListAsync(); + Assert.Equal(2, results.Count); + var firstRecord = Assert.Single(results, r => r.Key.Equals(key1)); + Assert.Equal(8, firstRecord.Int); + var secondRecord = Assert.Single(results, r => r.Key.Equals(key2)); + Assert.Equal(9, secondRecord.Int); + } /////////////////////// // Test dynamic mapping @@ -62,41 +86,120 @@ protected virtual async Task Test(TKey key1, TKey key2) await collection.DeleteAsync([key1, key2]); await fixture.TestStore.WaitForDataAsync(collection, recordCount: 0); - using var dynamicCollection = fixture.CreateDynamicCollection(); - await dynamicCollection.EnsureCollectionExistsAsync(); + using (var dynamicCollection = fixture.CreateDynamicCollection(withAutoGeneration: false)) + { + await dynamicCollection.EnsureCollectionExistsAsync(); + + var dynamicRecord = new Dictionary + { + [nameof(Record.Key)] = key1, + [nameof(Record.Int)] = 8, + [nameof(Record.Vector)] = new ReadOnlyMemory([1, 2, 3]) + }; + var nextDynamicRecord = new Dictionary + { + [nameof(Record.Key)] = key2, + [nameof(Record.Int)] = 9, + [nameof(Record.Vector)] = new ReadOnlyMemory([3, 2, 1]) + }; - var dynamicRecord = new Dictionary + await dynamicCollection.UpsertAsync(dynamicRecord); + // Exercise multi-record plus updating existing record + await dynamicCollection.UpsertAsync([dynamicRecord, nextDynamicRecord]); + await fixture.TestStore.WaitForDataAsync(dynamicCollection, recordCount: 2); + + // Single record get + var dynamicResult = await dynamicCollection.GetAsync(key1); + Assert.NotNull(dynamicResult); + Assert.IsType(dynamicResult[nameof(Record.Key)]); + Assert.Equal(key1, (TKey)dynamicResult[nameof(Record.Key)]!); + Assert.Equal(8, dynamicResult[nameof(Record.Int)]); + + // Multiple record get + // Also ensures that the second record - with the default key value - got properly inserted and did not trigger auto-generation + // (as we haven't configured it). + var dynamicResults = await dynamicCollection.GetAsync([key1, key2]).ToListAsync(); + Assert.Equal(2, dynamicResults.Count); + var firstDynamicRecord = Assert.Single(dynamicResults, r => r[nameof(Record.Key)]!.Equals(key1)); + Assert.IsType(firstDynamicRecord[nameof(Record.Key)]); + Assert.Equal(8, firstDynamicRecord[nameof(Record.Int)]); + var secondDynamicRecord = Assert.Single(dynamicResults, r => r[nameof(Record.Key)]!.Equals(key2)); + Assert.IsType(secondDynamicRecord[nameof(Record.Key)]); + Assert.Equal(9, secondDynamicRecord[nameof(Record.Int)]); + } + + if (supportsAutoGeneration) { - [nameof(Record.Key)] = key1, - [nameof(Record.Int)] = 8, - [nameof(Record.Vector)] = new ReadOnlyMemory([1, 2, 3]) - }; - var nextDynamicRecord = new Dictionary + // Above we tested with a collection where auto-generation isn't enabled - including with the default key value, + // which would have triggered auto-generation if it was enabled. + // Now, drop and recreate the collection with auto-generation enabled, and test that it works. + await collection.EnsureCollectionDeletedAsync(); + + // Pass null to test the provider's default behavior, which should be to enable auto-generation. + using var collectionWithAutoGeneration = fixture.CreateCollection(withAutoGeneration: null); + await collectionWithAutoGeneration.EnsureCollectionExistsAsync(); + + var record = new Record + { + Key = key1, + Int = 8, + Vector = new ReadOnlyMemory([1, 2, 3]) + }; + + var recordWithDefaultValueKey1 = new Record + { + Key = key2, + Int = 9, + Vector = new ReadOnlyMemory([3, 2, 1]) + }; + + var recordWithDefaultValueKey2 = new Record + { + Key = key2, + Int = 10, + Vector = new ReadOnlyMemory([3, 2, 1]) + }; + + var recordWithDefaultValueKey3 = new Record + { + Key = key2, + Int = 11, + Vector = new ReadOnlyMemory([3, 2, 1]) + }; + + // recordWithDefaultValueKey1 gets inserted alone, exercising single-record upsert with auto-generation. + await collectionWithAutoGeneration.UpsertAsync(recordWithDefaultValueKey1); + Assert.NotEqual(recordWithDefaultValueKey1.Key, key2); + var preUpdateGeneratedKey = recordWithDefaultValueKey1.Key; + recordWithDefaultValueKey1.Int = 99; + + // recordWithDefaultValueKey1 gets upserted, exercising update instead of insert. + // recordWithDefaultValueKey2 and 3 get inserted, exercising multi-record upsert with auto-generation; we insert two records to make + // sure the correct key gets injected back into each record. + // Finally, record gets inserted with a non-generated key, to make sure auto-generation doesn't kick in for non-CLR-default keys. + await collectionWithAutoGeneration.UpsertAsync([recordWithDefaultValueKey1, recordWithDefaultValueKey2, recordWithDefaultValueKey3, record]); + await fixture.TestStore.WaitForDataAsync(collectionWithAutoGeneration, recordCount: 4); + + Assert.Equal(recordWithDefaultValueKey1.Key, preUpdateGeneratedKey); + Assert.Equal(99, recordWithDefaultValueKey1.Int); + Assert.NotEqual(recordWithDefaultValueKey2.Key, key2); + Assert.NotEqual(recordWithDefaultValueKey3.Key, key2); + Assert.NotEqual(recordWithDefaultValueKey2.Key, recordWithDefaultValueKey1.Key!); + Assert.NotEqual(recordWithDefaultValueKey3.Key, recordWithDefaultValueKey1.Key!); + Assert.NotEqual(recordWithDefaultValueKey3.Key, recordWithDefaultValueKey2.Key!); + Assert.Equal(record.Key, key1); + + var results = await collectionWithAutoGeneration.GetAsync([key1, recordWithDefaultValueKey1.Key, recordWithDefaultValueKey2.Key, recordWithDefaultValueKey3.Key]).ToListAsync(); + Assert.Single(results, r => r.Key.Equals(recordWithDefaultValueKey1.Key)); + Assert.Single(results, r => r.Key.Equals(recordWithDefaultValueKey2.Key)); + Assert.Single(results, r => r.Key.Equals(recordWithDefaultValueKey3.Key)); + Assert.Single(results, r => r.Key.Equals(key1)); + } + else { - [nameof(Record.Key)] = key2, - [nameof(Record.Int)] = 9, - [nameof(Record.Vector)] = new ReadOnlyMemory([3, 2, 1]) - }; - - await dynamicCollection.UpsertAsync(dynamicRecord); - await dynamicCollection.UpsertAsync([dynamicRecord, nextDynamicRecord]); - await fixture.TestStore.WaitForDataAsync(dynamicCollection, recordCount: 2); - - var dynamicResult = await dynamicCollection.GetAsync(key1); - - Assert.NotNull(dynamicResult); - Assert.IsType(dynamicResult[nameof(Record.Key)]); - Assert.Equal(key1, (TKey)dynamicResult[nameof(Record.Key)]!); - Assert.Equal(8, dynamicResult[nameof(Record.Int)]); - - var dynamicResults = await dynamicCollection.GetAsync([key1, key2]).ToListAsync(); - Assert.Equal(2, dynamicResults.Count); - var firstDynamicRecord = Assert.Single(dynamicResults, r => r[nameof(Record.Key)]!.Equals(key1)); - Assert.IsType(firstDynamicRecord[nameof(Record.Key)]); - Assert.Equal(8, firstDynamicRecord[nameof(Record.Int)]); - var secondDynamicRecord = Assert.Single(dynamicResults, r => r[nameof(Record.Key)]!.Equals(key2)); - Assert.IsType(secondDynamicRecord[nameof(Record.Key)]); - Assert.Equal(9, secondDynamicRecord[nameof(Record.Int)]); + // Auto-generation is not supported for this type; ensure that model validation throws. + Assert.Throws(() => fixture.CreateCollection(withAutoGeneration: true)); + } } public abstract class Fixture : VectorStoreFixture @@ -104,21 +207,21 @@ public abstract class Fixture : VectorStoreFixture protected virtual string CollectionNameBase => nameof(KeyTypeTests); public virtual string CollectionName => this.TestStore.AdjustCollectionName(this.CollectionNameBase); - public virtual VectorStoreCollection> CreateCollection() + public virtual VectorStoreCollection> CreateCollection(bool? withAutoGeneration) where TKey : notnull - => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName, this.CreateRecordDefinition()); + => this.TestStore.DefaultVectorStore.GetCollection>(this.CollectionName, this.CreateRecordDefinition(withAutoGeneration)); - public virtual VectorStoreCollection> CreateDynamicCollection() + public virtual VectorStoreCollection> CreateDynamicCollection(bool withAutoGeneration) where TKey : notnull - => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition()); + => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition(withAutoGeneration)); - public virtual VectorStoreCollectionDefinition CreateRecordDefinition() + public virtual VectorStoreCollectionDefinition CreateRecordDefinition(bool? withAutoGeneration) where TKey : notnull => new() { Properties = [ - new VectorStoreKeyProperty("Key", typeof(TKey)), + new VectorStoreKeyProperty("Key", typeof(TKey)) { IsAutoGenerated = withAutoGeneration }, new VectorStoreDataProperty("Int", typeof(int)), new VectorStoreVectorProperty("Vector", typeof(ReadOnlyMemory), dimensions: 3) { diff --git a/dotnet/test/VectorData/VectorData.UnitTests/CollectionModelBuilderTests.cs b/dotnet/test/VectorData/VectorData.UnitTests/CollectionModelBuilderTests.cs index 2971a8811b4d..0bc26fea1ef5 100644 --- a/dotnet/test/VectorData/VectorData.UnitTests/CollectionModelBuilderTests.cs +++ b/dotnet/test/VectorData/VectorData.UnitTests/CollectionModelBuilderTests.cs @@ -373,11 +373,15 @@ private sealed class CustomModelBuilder(CollectionModelBuildingOptions? options RequiresAtLeastOneVector = false }; - protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes) + protected override void ValidateKeyProperty(KeyPropertyModel keyProperty) { - supportedTypes = "string, int"; + var type = keyProperty.Type; - return type == typeof(string) || type == typeof(int); + if (type != typeof(string) && type != typeof(int)) + { + throw new NotSupportedException( + $"Property '{keyProperty.ModelName}' has unsupported type '{type.Name}'. Key properties must be one of the supported types: string, int."); + } } protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes)