.Net: Optimize DuckDb GetBatchAsync to retrieve multiple entries with a single query

dotnet/src/Connectors/Connectors.Memory.DuckDB/Database.cs

@@ -183,6 +183,53 @@ ORDER BY score DESC
         return null;
     }
 
+    public async IAsyncEnumerable<DatabaseEntry> ReadAllAsync(DuckDBConnection conn,
+    string collectionName,
+    string[] keys,
+    [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        if (keys.Length == 0)
+        {
+            yield break;
+        }
+
+        using var cmd = conn.CreateCommand();
+        var keyPlaceholders = string.Join(", ", keys.Select((k) => "?"));
+
+#pragma warning disable CA2100 // Review SQL queries for security vulnerabilities
+        cmd.CommandText = $@"
+         SELECT key, metadata, timestamp, embedding 
+         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 });
+        }
+
+        using var dataReader = await cmd.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);
+
+        while (await dataReader.ReadAsync(cancellationToken).ConfigureAwait(false))
+        {
+            string key = dataReader.GetFieldValue<string>("key");
+            string metadata = dataReader.GetFieldValue<string>("metadata");
+            string timestamp = dataReader.GetFieldValue<string>("timestamp");
+            var embeddingFromSearch = dataReader.GetFieldValue<List<float>>("embedding").ToArray();
+
+            yield return new DatabaseEntry
+            {
+                Key = key,
+                MetadataString = metadata,
+                Embedding = embeddingFromSearch,
+                Timestamp = timestamp
+            };
+        }
+    }
+
     public async Task DeleteCollectionAsync(DuckDBConnection conn, string collectionName, CancellationToken cancellationToken = default)
     {
         using var cmd = conn.CreateCommand();

dotnet/src/Connectors/Connectors.Memory.DuckDB/DuckDBMemoryStore.cs

@@ -144,17 +144,9 @@ public async IAsyncEnumerable<string> UpsertBatchAsync(string collectionName, IE
     public async IAsyncEnumerable<MemoryRecord> GetBatchAsync(string collectionName, IEnumerable<string> keys, bool withEmbeddings = false,
         [EnumeratorCancellation] CancellationToken cancellationToken = default)
     {
-        foreach (var key in keys)
+        await foreach (var memoryRecord in this.InternalGetBatchAsync(this._dbConnection, collectionName, keys.ToArray(), withEmbeddings, cancellationToken).ConfigureAwait(false))
         {
-            var result = await this.InternalGetAsync(this._dbConnection, collectionName, key, withEmbeddings, cancellationToken).ConfigureAwait(false);
-            if (result is not null)
-            {
-                yield return result;
-            }
-            else
-            {
-                yield break;
-            }
+            yield return memoryRecord;
         }
     }
 
@@ -320,5 +312,30 @@ await this._dbConnector.UpdateOrInsertAsync(conn: connection,
             ParseTimestamp(entry.Value.Timestamp));
     }
 
+    private async IAsyncEnumerable<MemoryRecord> InternalGetBatchAsync(
+    DuckDBConnection connection,
+    string collectionName,
+    string[] keys,
+    bool withEmbedding,
+    [EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        await foreach (var entry in this._dbConnector.ReadAllAsync(connection, collectionName, keys, cancellationToken).ConfigureAwait(false))
+        {
+            var memoryRecord = withEmbedding
+                ? MemoryRecord.FromJsonMetadata(
+                    json: entry.MetadataString,
+                    entry.Embedding,
+                    entry.Key,
+                    ParseTimestamp(entry.Timestamp))
+                : MemoryRecord.FromJsonMetadata(
+                    json: entry.MetadataString,
+                    ReadOnlyMemory<float>.Empty,
+                    entry.Key,
+                    ParseTimestamp(entry.Timestamp));
+
+            yield return memoryRecord;
+        }
+    }
+
     #endregion
 }