-
Notifications
You must be signed in to change notification settings - Fork 5
sqlitevec: use DELETE by key instead of IN for virtual table deletes #53
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
rossdonald
wants to merge
2
commits into
CommunityToolkit:main
Choose a base branch
from
rossdonald:sqlite-batch-delete-with-key
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+68
−36
Open
Changes from all commits
Commits
Show all changes
2 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 /> | ||
|
|
@@ -617,16 +611,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, cancellationToken).ConfigureAwait(false); | ||
|
|
||
| using var vectorInsertCommand = SqliteCommandBuilder.BuildInsertCommand( | ||
| connection, | ||
|
|
@@ -644,36 +629,47 @@ await connection.ExecuteWithErrorHandlingAsync( | |
| } | ||
| } | ||
|
|
||
| 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>(); | ||
|
|
||
| if (_vectorPropertiesExist) | ||
| { | ||
| using var vectorCommand = SqliteCommandBuilder.BuildDeleteCommand( | ||
| connection, | ||
| _vectorTableName, | ||
| [condition]); | ||
|
|
||
| tasks.Add(connection.ExecuteWithErrorHandlingAsync( | ||
| _collectionMetadata, | ||
| "VectorDelete", | ||
| () => vectorCommand.ExecuteNonQueryAsync(cancellationToken), | ||
| cancellationToken)); | ||
| await DeleteVectorRowsAsync(connection, keys, 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)]); | ||
|
|
||
| tasks.Add(connection.ExecuteWithErrorHandlingAsync( | ||
| await connection.ExecuteWithErrorHandlingAsync( | ||
| _collectionMetadata, | ||
| "DataDelete", | ||
| () => dataCommand.ExecuteNonQueryAsync(cancellationToken), | ||
| cancellationToken)); | ||
| cancellationToken).ConfigureAwait(false); | ||
| } | ||
|
|
||
| private async Task DeleteVectorRowsAsync(SqliteConnection connection, IEnumerable<object> keys, 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); | ||
|
|
||
| var keyParameter = vectorDeleteCommand.Parameters[SqliteCommandBuilder.KeyParameterName]; | ||
|
|
||
| return Task.WhenAll(tasks); | ||
| foreach (var key in keys) | ||
| { | ||
| keyParameter.Value = key; | ||
|
|
||
| await connection.ExecuteWithErrorHandlingAsync( | ||
|
Comment on lines
+663
to
+667
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @rossdonald It sounds reasonable to address it. You could re-use the benchmarks copilot has created for me to measure the difference: Details// Benchmarks for https://github.com/CommunityToolkit/AI/issues/52
// Compares SqliteVec batch delete/upsert cost on the vec0 virtual table.
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using BenchmarkDotNet.Attributes;
using BenchmarkDotNet.Configs;
using BenchmarkDotNet.Running;
using CommunityToolkit.VectorData.SqliteVec;
using Microsoft.Extensions.VectorData;
namespace SqliteVecBench;
public sealed class Record
{
[VectorStoreKey(StorageName = "chunk_id")]
public string ChunkId { get; set; }
[VectorStoreData]
public string Text { get; set; }
[VectorStoreVector(Dimensions: 256, DistanceFunction = DistanceFunction.CosineDistance)]
public ReadOnlyMemory<float> Embedding { get; set; }
}
[MemoryDiagnoser(displayGenColumns: false)]
public class SqliteVecDeleteBenchmarks
{
private const int BatchSize = 200;
private string _dbPath;
private SqliteCollection<string, Record> _collection;
private List<string> _missingKeys;
private List<string> _existingKeys;
private List<Record> _existingRecords;
[Params(10_000, 100_000)]
public int RowCount { get; set; }
private static ReadOnlyMemory<float> CreateVector(Random random)
{
float[] values = new float[256];
for (int i = 0; i < values.Length; i++)
{
values[i] = (float)random.NextDouble();
}
return new ReadOnlyMemory<float>(values);
}
[GlobalSetup]
public async Task SetupAsync()
{
_dbPath = Path.Combine(Path.GetTempPath(), $"sqlitevec-bench-{RowCount}-{Guid.NewGuid():N}.db");
_collection = new SqliteCollection<string, Record>($"Data Source={_dbPath}", "vec_chunks");
await _collection.EnsureCollectionExistsAsync();
Random random = new Random(42);
List<Record> batch = new List<Record>(1000);
for (int i = 0; i < RowCount; i++)
{
batch.Add(new Record { ChunkId = $"key-{i}", Text = "text", Embedding = CreateVector(random) });
if (batch.Count == 1000)
{
await _collection.UpsertAsync(batch);
batch.Clear();
}
}
if (batch.Count > 0)
{
await _collection.UpsertAsync(batch);
}
_missingKeys = Enumerable.Range(0, BatchSize).Select(i => $"missing-{i}").ToList();
// Keys spread across the table, from the "middle" of the key space.
_existingKeys = Enumerable.Range(0, BatchSize).Select(i => $"key-{i * (RowCount / BatchSize)}").ToList();
_existingRecords = _existingKeys
.Select(k => new Record { ChunkId = k, Text = "text", Embedding = CreateVector(random) })
.ToList();
}
[GlobalCleanup]
public void Cleanup()
{
_collection?.Dispose();
Microsoft.Data.Sqlite.SqliteConnection.ClearAllPools();
if (_dbPath is not null && File.Exists(_dbPath))
{
File.Delete(_dbPath);
}
}
// Delete a batch of keys that are not in the table: no rows are removed, so the
// benchmark is idempotent and measures the lookup cost only.
[Benchmark]
public Task DeleteBatch_MissingKeys() => _collection.DeleteAsync(_missingKeys);
// Single key delete for a key that is not present.
[Benchmark]
public Task DeleteSingle_MissingKey() => _collection.DeleteAsync("missing-0");
// Upsert of records that already exist: internally deletes the vector rows of the
// batch and re-inserts them, so the table size stays constant.
[Benchmark]
public Task UpsertBatch_ExistingRecords() => _collection.UpsertAsync(_existingRecords);
}
public static class Program
{
public static void Main(string[] args)
=> BenchmarkSwitcher.FromAssembly(typeof(Program).Assembly).Run(args);
} |
||
| _collectionMetadata, | ||
| "VectorDelete", | ||
| () => vectorDeleteCommand.ExecuteNonQueryAsync(cancellationToken), | ||
| cancellationToken).ConfigureAwait(false); | ||
| } | ||
| } | ||
|
|
||
| /// <summary> | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@rossdonald please bump the version to 1.0.2-preview here:
AI/MEVD/src/SqliteVec/SqliteVec.csproj
Line 4 in 215a5ba
(I am going to release a new version to nuget.org as soon as this PR gets merged)