diff --git a/Directory.Packages.props b/Directory.Packages.props index 2e4180a..4c5f90e 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -25,7 +25,7 @@ - + diff --git a/MEVD/src/SqliteVec/SqliteCollection.cs b/MEVD/src/SqliteVec/SqliteCollection.cs index 847d534..893dfc1 100644 --- a/MEVD/src/SqliteVec/SqliteCollection.cs +++ b/MEVD/src/SqliteVec/SqliteCollection.cs @@ -347,9 +347,7 @@ public override async Task DeleteAsync(TKey key, CancellationToken cancellationT using var connection = await GetConnectionAsync(cancellationToken).ConfigureAwait(false); - var condition = new SqliteWhereEqualsCondition(_keyStorageName, key); - - await InternalDeleteBatchAsync(connection, condition, cancellationToken).ConfigureAwait(false); + await InternalDeleteBatchAsync(connection, [key], cancellationToken).ConfigureAwait(false); } /// @@ -364,11 +362,7 @@ public override async Task DeleteAsync(IEnumerable keys, CancellationToken using var connection = await GetConnectionAsync(cancellationToken).ConfigureAwait(false); - var condition = new SqliteWhereInCondition( - _keyStorageName, - keysList); - - await InternalDeleteBatchAsync(connection, condition, cancellationToken).ConfigureAwait(false); + await InternalDeleteBatchAsync(connection, keysList, cancellationToken).ConfigureAwait(false); } /// @@ -568,6 +562,8 @@ private async Task DoUpsertAsync(IEnumerable records, CancellationToken using var connection = await GetConnectionAsync(cancellationToken).ConfigureAwait(false); + using var transaction = connection.BeginTransaction(); + using var dataCommand = SqliteCommandBuilder.BuildInsertCommand( connection, _dataTableName, @@ -576,6 +572,7 @@ private async Task DoUpsertAsync(IEnumerable records, CancellationToken generatedEmbeddings, data: true, replaceIfExists: true); + dataCommand.Transaction = transaction; using (var reader = await connection.ExecuteWithErrorHandlingAsync( _collectionMetadata, @@ -617,16 +614,7 @@ private async Task DoUpsertAsync(IEnumerable records, CancellationToken // Deleting vector records first since current version of vector search extension // doesn't support Upsert operation, only Delete/Insert. - using var vectorDeleteCommand = SqliteCommandBuilder.BuildDeleteCommand( - connection, - _vectorTableName, - [new SqliteWhereInCondition(_keyStorageName, keys)]); - - await connection.ExecuteWithErrorHandlingAsync( - _collectionMetadata, - "VectorDelete", - () => vectorDeleteCommand.ExecuteNonQueryAsync(cancellationToken), - cancellationToken).ConfigureAwait(false); + await DeleteVectorRowsAsync(connection, keys, transaction, cancellationToken).ConfigureAwait(false); using var vectorInsertCommand = SqliteCommandBuilder.BuildInsertCommand( connection, @@ -635,6 +623,7 @@ await connection.ExecuteWithErrorHandlingAsync( recordsList, generatedEmbeddings, data: false); + vectorInsertCommand.Transaction = transaction; await connection.ExecuteWithErrorHandlingAsync( _collectionMetadata, @@ -642,38 +631,57 @@ await connection.ExecuteWithErrorHandlingAsync( () => vectorInsertCommand.ExecuteNonQueryAsync(cancellationToken), cancellationToken).ConfigureAwait(false); } + + transaction.Commit(); } - private Task InternalDeleteBatchAsync(SqliteConnection connection, SqliteWhereCondition condition, CancellationToken cancellationToken) + private async Task InternalDeleteBatchAsync(SqliteConnection connection, List keys, CancellationToken cancellationToken) { - var tasks = new List(); + using var transaction = connection.BeginTransaction(); if (_vectorPropertiesExist) { - using var vectorCommand = SqliteCommandBuilder.BuildDeleteCommand( - connection, - _vectorTableName, - [condition]); - - tasks.Add(connection.ExecuteWithErrorHandlingAsync( - _collectionMetadata, - "VectorDelete", - () => vectorCommand.ExecuteNonQueryAsync(cancellationToken), - cancellationToken)); + await DeleteVectorRowsAsync(connection, keys, transaction, cancellationToken).ConfigureAwait(false); } + // The data table is a regular table with an indexed primary key, so DELETE using IN is efficient. using var dataCommand = SqliteCommandBuilder.BuildDeleteCommand( connection, _dataTableName, - [condition]); + [new SqliteWhereInCondition(_keyStorageName, keys)]); + dataCommand.Transaction = transaction; - tasks.Add(connection.ExecuteWithErrorHandlingAsync( + await connection.ExecuteWithErrorHandlingAsync( _collectionMetadata, "DataDelete", () => dataCommand.ExecuteNonQueryAsync(cancellationToken), - cancellationToken)); + cancellationToken).ConfigureAwait(false); + + transaction.Commit(); + } + + private async Task DeleteVectorRowsAsync(SqliteConnection connection, IEnumerable keys, SqliteTransaction transaction, CancellationToken cancellationToken) + { + // One DELETE per key because the vec0 virtual table cannot use an IN-list, so a single + // batched DELETE would scan the whole table instead of using the primary key. + using var vectorDeleteCommand = SqliteCommandBuilder.BuildDeleteByKeyCommand( + connection, + _vectorTableName, + _keyStorageName); + vectorDeleteCommand.Transaction = transaction; + + var keyParameter = vectorDeleteCommand.Parameters[SqliteCommandBuilder.KeyParameterName]; - return Task.WhenAll(tasks); + foreach (var key in keys) + { + keyParameter.Value = key; + + await connection.ExecuteWithErrorHandlingAsync( + _collectionMetadata, + "VectorDelete", + () => vectorDeleteCommand.ExecuteNonQueryAsync(cancellationToken), + cancellationToken).ConfigureAwait(false); + } } /// diff --git a/MEVD/src/SqliteVec/SqliteCommandBuilder.cs b/MEVD/src/SqliteVec/SqliteCommandBuilder.cs index 29eebf5..5e0817e 100644 --- a/MEVD/src/SqliteVec/SqliteCommandBuilder.cs +++ b/MEVD/src/SqliteVec/SqliteCommandBuilder.cs @@ -20,6 +20,7 @@ namespace CommunityToolkit.VectorData.SqliteVec; internal static class SqliteCommandBuilder { internal const string DistancePropertyName = "distance"; + internal const string KeyParameterName = "@key"; public static DbCommand BuildTableCountCommand(SqliteConnection connection, string tableName) { @@ -387,6 +388,24 @@ public static DbCommand BuildDeleteCommand( return command; } + public static DbCommand BuildDeleteByKeyCommand( + SqliteConnection connection, + string tableName, + string keyColumnName) + { + var command = connection.CreateCommand(); + + command.CommandText = new StringBuilder() + .Append("DELETE FROM ").AppendIdentifier(tableName) + .Append(" WHERE ").AppendIdentifier(keyColumnName) + .Append(" = ").Append(KeyParameterName) + .ToString(); + + command.Parameters.Add(new SqliteParameter { ParameterName = KeyParameterName }); + + return command; + } + /// /// Appends a properly quoted and escaped SQLite identifier to the StringBuilder. /// In SQLite, identifiers are quoted with double quotes, and embedded double quotes are escaped by doubling them. diff --git a/MEVD/src/SqliteVec/SqliteVec.csproj b/MEVD/src/SqliteVec/SqliteVec.csproj index 8661550..40a10b1 100644 --- a/MEVD/src/SqliteVec/SqliteVec.csproj +++ b/MEVD/src/SqliteVec/SqliteVec.csproj @@ -1,7 +1,7 @@  - 1.0.1-preview + 1.0.2-preview CommunityToolkit.VectorData.SqliteVec $(AssemblyName) net10.0;net8.0;netstandard2.0;net462 diff --git a/MEVD/test/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs b/MEVD/test/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs index 93044ca..9f209ca 100644 --- a/MEVD/test/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs +++ b/MEVD/test/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs @@ -376,6 +376,23 @@ public void ItBuildsDeleteCommand() Assert.Equal(30, command.Parameters[3].Value); } + [Fact] + public void ItBuildsDeleteByKeyCommand() + { + // Arrange + const string TableName = "TestTable"; + const string KeyName = "Id"; + + // Act + var command = SqliteCommandBuilder.BuildDeleteByKeyCommand(this._connection, TableName, KeyName); + + // Assert + Assert.Equal("DELETE FROM \"TestTable\" WHERE \"Id\" = " + SqliteCommandBuilder.KeyParameterName, command.CommandText); + + Assert.Equal(SqliteCommandBuilder.KeyParameterName, command.Parameters[0].ParameterName); + Assert.Null(command.Parameters[0].Value); + } + public void Dispose() { this._command.Dispose();