.Net: Optimize Redis GetBatchAsync for Performance Improvement

dotnet/src/Connectors/Connectors.Memory.Redis/RedisMemoryStore.cs

@@ -141,13 +141,9 @@ public async Task DeleteCollectionAsync(string collectionName, CancellationToken
     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 result in this.InternalGetBatchAsync(collectionName, keys.ToArray(), withEmbeddings, cancellationToken).ConfigureAwait(false))
         {
-            var result = await this.InternalGetAsync(collectionName, key, withEmbeddings, cancellationToken).ConfigureAwait(false);
-            if (result is not null)
-            {
-                yield return result;
-            }
+            yield return result;
         }
     }
 
@@ -359,6 +355,49 @@ private static RedisKey GetRedisKey(string collectionName, string key)
             timestamp: ParseTimestamp((long?)hashEntries.FirstOrDefault(x => x.Name == "timestamp").Value));
     }
 
+    private async IAsyncEnumerable<MemoryRecord> InternalGetBatchAsync(string collectionName, string[] keys, bool withEmbeddings, [EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        var batch = this._database.CreateBatch();
+        var tasks = keys.Select(async key =>
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            HashEntry[] hashEntries = await this._database.HashGetAllAsync(GetRedisKey(collectionName, key), CommandFlags.None).ConfigureAwait(false);
+
+            if (hashEntries.Length == 0)
+            {
+                return null;
+            }
+
+            if (withEmbeddings)
+            {
+                RedisValue embedding = hashEntries.FirstOrDefault(x => x.Name == "embedding").Value;
+                return MemoryRecord.FromJsonMetadata(
+                    json: hashEntries.FirstOrDefault(x => x.Name == "metadata").Value!,
+                    embedding: embedding.HasValue ? MemoryMarshal.Cast<byte, float>((byte[])embedding!).ToArray() : ReadOnlyMemory<float>.Empty,
+                    key: hashEntries.FirstOrDefault(x => x.Name == "key").Value,
+                    timestamp: ParseTimestamp((long?)hashEntries.FirstOrDefault(x => x.Name == "timestamp").Value));
+            }
+
+            return MemoryRecord.FromJsonMetadata(
+                json: hashEntries.FirstOrDefault(x => x.Name == "metadata").Value!,
+                embedding: ReadOnlyMemory<float>.Empty,
+                key: hashEntries.FirstOrDefault(x => x.Name == "key").Value,
+                timestamp: ParseTimestamp((long?)hashEntries.FirstOrDefault(x => x.Name == "timestamp").Value));
+        });
+
+        batch.Execute();
+        var results = await Task.WhenAll(tasks).ConfigureAwait(false);
+
+        foreach (var result in results)
+        {
+            if (result is not null)
+            {
+                yield return result;
+            }
+        }
+    }
+
     private double GetSimilarity(Document document)
     {
         RedisValue vectorScoreValue = document["vector_score"];