Skip to content

Commit

Permalink
.Net: DuckDBMemoryStore RemoveBatchAsync method performance improveme…
Browse files Browse the repository at this point in the history
…nt (#7252)

### Motivation and Context

To optimize performance by reducing the number of database calls during
batch deletions in DuckDBMemoryStore.

### Description

The updated **RemoveBatchAsync** method now executes a single SQL DELETE
query to remove multiple keys at once, replacing the previous approach
of executing individual delete queries for each key.

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [ ] The code builds clean without any errors or warnings
- [ ] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [ ] All unit tests pass, and existing test is okay for this change
- [ ] I didn't break anyone 😄
  • Loading branch information
atiq-bs23 authored Jul 15, 2024
1 parent b05478d commit 6cfd2ec
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 3 deletions.
25 changes: 25 additions & 0 deletions dotnet/src/Connectors/Connectors.Memory.DuckDB/Database.cs
Original file line number Diff line number Diff line change
Expand Up @@ -204,4 +204,29 @@ DELETE FROM {TableName}
cmd.Parameters.Add(new DuckDBParameter(nameof(key), key));
await cmd.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
}

public async Task DeleteBatchAsync(DuckDBConnection conn, string collectionName, string[] keys, CancellationToken cancellationToken = default)
{
if (keys.Length == 0)
{
return;
}

using var cmd = conn.CreateCommand();
var keyPlaceholders = string.Join(", ", keys.Select((k) => "?"));

#pragma warning disable CA2100 // Review SQL queries for security vulnerabilities
cmd.CommandText = $@"
DELETE FROM {TableName}
WHERE collection=?
AND key IN ({keyPlaceholders});";
#pragma warning restore CA2100 // Review SQL queries for security vulnerabilities
cmd.Parameters.Add(new DuckDBParameter { Value = collectionName });
// Add the key parameters
foreach (var key in keys)
{
cmd.Parameters.Add(new DuckDBParameter { Value = key });
}
await cmd.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ public async Task RemoveAsync(string collectionName, string key, CancellationTok
/// <inheritdoc/>
public async Task RemoveBatchAsync(string collectionName, IEnumerable<string> keys, CancellationToken cancellationToken = default)
{
await Task.WhenAll(keys.Select(k => this._dbConnector.DeleteAsync(this._dbConnection, collectionName, k, cancellationToken))).ConfigureAwait(false);
await this._dbConnector.DeleteBatchAsync(this._dbConnection, collectionName, keys.ToArray(), cancellationToken).ConfigureAwait(false);
}

/// <inheritdoc/>
Expand Down Expand Up @@ -271,8 +271,7 @@ private DuckDBMemoryStore(DuckDBConnection connection, int? vectorSize = null)

private static DateTimeOffset? ParseTimestamp(string? str)
{
if (!string.IsNullOrEmpty(str)
&& DateTimeOffset.TryParse(str, CultureInfo.InvariantCulture, DateTimeStyles.AssumeUniversal, out DateTimeOffset timestamp))
if (DateTimeOffset.TryParse(str, CultureInfo.InvariantCulture, DateTimeStyles.AssumeUniversal, out DateTimeOffset timestamp))
{
return timestamp;
}
Expand Down

0 comments on commit 6cfd2ec

Please sign in to comment.