Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Directory.Packages.props
Original file line number Diff line number Diff line change
Expand Up @@ -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">
Expand Down
66 changes: 31 additions & 35 deletions MEVD/src/SqliteVec/SqliteCollection.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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 />
Expand All @@ -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 />
Expand Down Expand Up @@ -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);

Copy link
Copy Markdown
Contributor

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:

<Version>1.0.1-preview</Version>

(I am going to release a new version to nuget.org as soon as this PR gets merged)


using var vectorInsertCommand = SqliteCommandBuilder.BuildInsertCommand(
connection,
Expand All @@ -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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The 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>
Expand Down
19 changes: 19 additions & 0 deletions MEVD/src/SqliteVec/SqliteCommandBuilder.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand Down Expand Up @@ -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.
Expand Down
17 changes: 17 additions & 0 deletions MEVD/test/SqliteVec.UnitTests/SqliteCommandBuilderTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading