sqlitevec: use DELETE by key instead of IN for virtual table deletes

Directory.Packages.props

@@ -25,7 +25,7 @@
   </ItemGroup>
 
   <ItemGroup Label="Build">
-    <PackageVersion Include="Microsoft.SourceLink.GitHub" Version="10.0.300" />
+    <PackageVersion Include="Microsoft.SourceLink.GitHub" Version="10.0.401" />
   </ItemGroup>
 
   <ItemGroup Label="Test">

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);
     }
 
     /// <inheritdoc />
@@ -364,11 +362,7 @@ public override async Task DeleteAsync(IEnumerable<TKey> 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);
     }
 
     /// <inheritdoc />
@@ -568,6 +562,8 @@ private async Task DoUpsertAsync(IEnumerable<TRecord> 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<TRecord> 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<TRecord> 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,45 +623,65 @@ await connection.ExecuteWithErrorHandlingAsync(
                 recordsList,
                 generatedEmbeddings,
                 data: false);
+            vectorInsertCommand.Transaction = transaction;
 
             await connection.ExecuteWithErrorHandlingAsync(
                 _collectionMetadata,
                 "VectorInsert",
                 () => vectorInsertCommand.ExecuteNonQueryAsync(cancellationToken),
                 cancellationToken).ConfigureAwait(false);
         }
+
+        transaction.Commit();
     }
 
-    private Task InternalDeleteBatchAsync(SqliteConnection connection, SqliteWhereCondition condition, CancellationToken cancellationToken)
+    private async Task InternalDeleteBatchAsync(SqliteConnection connection, List<object> keys, CancellationToken cancellationToken)
     {
-        var tasks = new List<Task>();
+        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<object> 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);
+        }
     }
 
     /// <summary>

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;
+    }
+
     /// <summary>
     /// 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.

MEVD/src/SqliteVec/SqliteVec.csproj

@@ -1,7 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
-    <Version>1.0.1-preview</Version>
+    <Version>1.0.2-preview</Version>
     <AssemblyName>CommunityToolkit.VectorData.SqliteVec</AssemblyName>
     <RootNamespace>$(AssemblyName)</RootNamespace>
     <TargetFrameworks>net10.0;net8.0;netstandard2.0;net462</TargetFrameworks>

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();