Implement QueryWithIterator/SearchWithIterator

Milvus.Client.Tests/QueryWithIteratorExtendedTests.cs

@@ -0,0 +1,179 @@
+using Xunit;
+
+namespace Milvus.Client.Tests;
+
+[Collection("Milvus")]
+public class QueryWithIteratorExtendedTests : IAsyncLifetime
+{
+    [Fact]
+    public async Task QueryWithIterator_string_pk_with_offset()
+    {
+        await Collection.DropAsync();
+        await Client.CreateCollectionAsync(
+            CollectionName,
+            new[]
+            {
+                FieldSchema.CreateVarchar("id", 64, isPrimaryKey: true),
+                FieldSchema.CreateVarchar("text", 256),
+                FieldSchema.CreateFloatVector("float_vector", 4),
+            });
+
+        var stringCollection = Client.GetCollection(CollectionName);
+
+        await stringCollection.InsertAsync(
+        [
+            FieldData.CreateVarChar("id", new[] { "a", "b", "c", "d", "e", "f", "g", "h" }),
+            FieldData.CreateVarChar("text", new[] { "text1", "text2", "text3", "text4", "text5", "text6", "text7", "text8" }),
+            FieldData.CreateFloatVector("float_vector", new ReadOnlyMemory<float>[]
+            {
+                new[] { 1f, 0f, 0f, 0f },
+                new[] { 2f, 0f, 0f, 0f },
+                new[] { 3f, 0f, 0f, 0f },
+                new[] { 4f, 0f, 0f, 0f },
+                new[] { 5f, 0f, 0f, 0f },
+                new[] { 6f, 0f, 0f, 0f },
+                new[] { 7f, 0f, 0f, 0f },
+                new[] { 8f, 0f, 0f, 0f }
+            })
+        ]);
+        await stringCollection.FlushAsync();
+
+        await stringCollection.CreateIndexAsync("float_vector", IndexType.Flat, SimilarityMetricType.L2);
+        await stringCollection.WaitForIndexBuildAsync("float_vector");
+        await stringCollection.LoadAsync();
+
+        var queryParams = new QueryParameters
+        {
+            Offset = 3,
+            Limit = 4,
+            OutputFields = { "id", "text" }
+        };
+
+        var iterator = stringCollection.QueryWithIteratorAsync(parameters: queryParams);
+
+        List<IReadOnlyList<FieldData>> results = new();
+        await foreach (var result in iterator)
+        {
+            results.Add(result);
+        }
+
+        var allIds = results
+            .SelectMany(r => ((FieldData<string>)r.First(f => f.FieldName == "id")).Data)
+            .OrderBy(x => x)
+            .ToList();
+
+        Assert.True(allIds.Count <= 4);
+        Assert.True(allIds.Count > 0);
+    }
+
+    [Fact]
+    public async Task QueryWithIterator_with_offset_and_expression()
+    {
+        await Collection.InsertAsync(
+        [
+            FieldData.Create("id", Enumerable.Range(1, 20).Select(i => (long)i).ToArray()),
+            FieldData.CreateVarChar("varchar", Enumerable.Range(1, 20).Select(i => $"text{i}").ToArray()),
+            FieldData.CreateFloatVector("float_vector",
+                Enumerable.Range(1, 20).Select(i => new ReadOnlyMemory<float>(new[] { i, 0f, 0f, 0f })).ToArray())
+        ]);
+        await Collection.FlushAsync();
+
+        await Collection.CreateIndexAsync("float_vector", IndexType.Flat, SimilarityMetricType.L2);
+        await Collection.WaitForIndexBuildAsync("float_vector");
+        await Collection.LoadAsync();
+
+        var queryParams = new QueryParameters
+        {
+            Offset = 3,
+            Limit = 5,
+            OutputFields = { "id" }
+        };
+
+        var iterator = Collection.QueryWithIteratorAsync(expression: "id >= 5 and id <= 15", parameters: queryParams, batchSize: 2);
+
+        List<IReadOnlyList<FieldData>> results = new();
+        await foreach (var result in iterator)
+        {
+            results.Add(result);
+        }
+
+        var allIds = results
+            .SelectMany(r => ((FieldData<long>)r.First(f => f.FieldName == "id")).Data)
+            .OrderBy(x => x)
+            .ToList();
+
+        Assert.True(allIds.Count <= 5);
+        Assert.True(allIds.Count > 0);
+        Assert.All(allIds, id => Assert.True(id >= 5 && id <= 15));
+    }
+
+    [Fact]
+    public async Task QueryWithIterator_large_offset()
+    {
+        await Collection.InsertAsync(
+        [
+            FieldData.Create("id", Enumerable.Range(1, 100).Select(i => (long)i).ToArray()),
+            FieldData.CreateVarChar("varchar", Enumerable.Range(1, 100).Select(i => $"text{i}").ToArray()),
+            FieldData.CreateFloatVector("float_vector",
+                Enumerable.Range(1, 100).Select(i => new ReadOnlyMemory<float>(new[] { i, 0f, 0f, 0f })).ToArray())
+        ]);
+        await Collection.FlushAsync();
+
+        await Collection.CreateIndexAsync("float_vector", IndexType.Flat, SimilarityMetricType.L2);
+        await Collection.WaitForIndexBuildAsync("float_vector");
+        await Collection.LoadAsync();
+
+        var queryParams = new QueryParameters
+        {
+            Offset = 90,
+            Limit = 10,
+            OutputFields = { "id" }
+        };
+
+        var iterator = Collection.QueryWithIteratorAsync(parameters: queryParams);
+
+        List<IReadOnlyList<FieldData>> results = new();
+        await foreach (var result in iterator)
+        {
+            results.Add(result);
+        }
+
+        var allIds = results
+            .SelectMany(r => ((FieldData<long>)r.First(f => f.FieldName == "id")).Data)
+            .ToList();
+
+        Assert.True(allIds.Count <= 10);
+        Assert.All(allIds, id => Assert.True(id > 90));
+    }
+
+    public async Task InitializeAsync()
+    {
+        await Collection.DropAsync();
+        await Client.CreateCollectionAsync(
+            CollectionName,
+            new[]
+            {
+                FieldSchema.Create<long>("id", isPrimaryKey: true),
+                FieldSchema.CreateVarchar("varchar", 256),
+                FieldSchema.CreateFloatVector("float_vector", 4),
+            });
+    }
+
+    public Task DisposeAsync()
+    {
+        Client.Dispose();
+        return Task.CompletedTask;
+    }
+
+    private const string CollectionName = nameof(QueryWithIteratorExtendedTests);
+
+    private readonly MilvusClient Client;
+
+    private MilvusCollection Collection { get; }
+
+    public QueryWithIteratorExtendedTests(MilvusFixture milvusFixture)
+    {
+        Client = milvusFixture.CreateClient();
+        Collection = Client.GetCollection(CollectionName);
+    }
+}

Milvus.Client.Tests/SearchQueryIteratorLongKeyTests.cs → Milvus.Client.Tests/QueryWithIteratorLongKeyTests.cs

@@ -3,14 +3,14 @@
 namespace Milvus.Client.Tests;
 
 [Collection("Milvus")]
-public class SearchQueryIteratorLongKeyTests : IClassFixture<SearchQueryIteratorLongKeyTests.DataCollectionFixture>,
+public class QueryWithIteratorLongKeyTests : IClassFixture<QueryWithIteratorLongKeyTests.DataCollectionFixture>,
                                                IAsyncLifetime
 {
-    private const string CollectionName = nameof(SearchQueryIteratorLongKeyTests);
+    private const string CollectionName = nameof(QueryWithIteratorLongKeyTests);
     private readonly DataCollectionFixture _dataCollectionFixture;
     private readonly MilvusClient Client;
 
-    public SearchQueryIteratorLongKeyTests(MilvusFixture milvusFixture, DataCollectionFixture dataCollectionFixture)
+    public QueryWithIteratorLongKeyTests(MilvusFixture milvusFixture, DataCollectionFixture dataCollectionFixture)
     {
         Client = milvusFixture.CreateClient();
         _dataCollectionFixture = dataCollectionFixture;
@@ -55,17 +55,89 @@ await Collection.InsertAsync(
     }
 
     [Fact]
-    public void QueryWithIterator_OffsetNotZero()
+    public async Task QueryWithIterator_WithOffset()
     {
+        var items = new List<Item>
+        {
+            new(1, new[] { 10f, 20f }),
+            new(2, new[] { 30f, 40f }),
+            new(3, new[] { 50f, 60f }),
+            new(4, new[] { 70f, 80f }),
+            new(5, new[] { 90f, 100f })
+        };
+
+        await Collection.InsertAsync(
+        [
+            FieldData.Create("id", items.Select(x => x.Id).ToArray()),
+            FieldData.CreateFloatVector("float_vector", items.Select(x => x.Vector).ToArray())
+        ]);
+
         var queryParameters = new QueryParameters
         {
-            Offset = 1
+            Offset = 2,
+            Limit = 3,
+            OutputFields = { "id", "float_vector" }
         };
 
         var iterator = Collection.QueryWithIteratorAsync(parameters: queryParameters);
 
-        var exception = Assert.ThrowsAsync<MilvusException>(async () => await iterator.GetAsyncEnumerator().MoveNextAsync());
-        Assert.NotNull(exception);
+        List<IReadOnlyList<FieldData>> results = new();
+        await foreach (var result in iterator)
+        {
+            results.Add(result);
+        }
+
+        var returnedItems = results.SelectMany(ExtractItems).OrderBy(x => x.Id).ToList();
+        Assert.True(returnedItems.Count <= 3);
+        Assert.True(returnedItems.Count > 0);
+        Assert.All(returnedItems, item => Assert.True(item.Id >= 3));
+    }
+
+    [Fact]
+    public async Task QueryWithIterator_WithOffsetAndExpression()
+    {
+        var items = new List<Item>
+        {
+            new(1, new[] { 10f, 20f }),
+            new(2, new[] { 30f, 40f }),
+            new(3, new[] { 50f, 60f }),
+            new(4, new[] { 70f, 80f }),
+            new(5, new[] { 90f, 100f }),
+            new(6, new[] { 110f, 120f }),
+            new(7, new[] { 130f, 140f }),
+            new(8, new[] { 150f, 160f }),
+            new(9, new[] { 170f, 180f }),
+            new(10, new[] { 190f, 200f })
+        };
+
+        await Collection.InsertAsync(
+        [
+            FieldData.Create("id", items.Select(x => x.Id).ToArray()),
+            FieldData.CreateFloatVector("float_vector", items.Select(x => x.Vector).ToArray())
+        ]);
+
+        var queryParameters = new QueryParameters
+        {
+            Offset = 2,
+            Limit = 4,
+            OutputFields = { "id", "float_vector" }
+        };
+
+        var iterator = Collection.QueryWithIteratorAsync(
+            expression: "id >= 3 and id <= 8",
+            parameters: queryParameters);
+
+        List<IReadOnlyList<FieldData>> results = new();
+        await foreach (var result in iterator)
+        {
+            results.Add(result);
+        }
+
+        var returnedItems = results.SelectMany(ExtractItems).OrderBy(x => x.Id).ToList();
+        Assert.True(returnedItems.Count <= 4);
+        Assert.True(returnedItems.Count > 0);
+        Assert.All(returnedItems, item => Assert.True(item.Id >= 3 && item.Id <= 8));
+        Assert.True(returnedItems[0].Id >= 5);
     }
 
     [Fact]

Milvus.Client.Tests/SearchQueryIteratorStringKeyTests.cs → Milvus.Client.Tests/QueryWithIteratorStringKeyTests.cs

@@ -3,15 +3,15 @@
 namespace Milvus.Client.Tests;
 
 [Collection("Milvus")]
-public class SearchQueryIteratorStringKeyTests : IClassFixture<SearchQueryIteratorStringKeyTests.DataCollectionFixture>,
+public class QueryWithIteratorStringKeyTests : IClassFixture<QueryWithIteratorStringKeyTests.DataCollectionFixture>,
                                                  IAsyncLifetime
 {
-    private const string CollectionName = nameof(SearchQueryIteratorStringKeyTests);
+    private const string CollectionName = nameof(QueryWithIteratorStringKeyTests);
 
     private readonly DataCollectionFixture _dataCollectionFixture;
     private readonly MilvusClient Client;
 
-    public SearchQueryIteratorStringKeyTests(MilvusFixture milvusFixture, DataCollectionFixture dataCollectionFixture)
+    public QueryWithIteratorStringKeyTests(MilvusFixture milvusFixture, DataCollectionFixture dataCollectionFixture)
     {
         Client = milvusFixture.CreateClient();
         _dataCollectionFixture = dataCollectionFixture;

Milvus.Client.Tests/SearchWithIteratorExtendedTests.cs

@@ -0,0 +1,87 @@
+using Xunit;
+
+namespace Milvus.Client.Tests;
+
+[Collection("Milvus")]
+public class SearchWithIteratorExtendedTests : IAsyncLifetime
+{
+    [Fact]
+    public async Task SearchWithIterator_binary_vector()
+    {
+        await Collection.DropAsync();
+        await Client.CreateCollectionAsync(
+            CollectionName,
+            new[]
+            {
+                FieldSchema.Create<long>("id", isPrimaryKey: true),
+                FieldSchema.CreateBinaryVector("binary_vector", 8),
+            });
+
+        var binaryCollection = Client.GetCollection(CollectionName);
+
+        await binaryCollection.InsertAsync(
+        [
+            FieldData.Create("id", new[] { 1L, 2L, 3L, 4L, 5L }),
+            FieldData.CreateBinaryVectors("binary_vector", new ReadOnlyMemory<byte>[]
+            {
+                new byte[] { 0x01 },
+                new byte[] { 0x02 },
+                new byte[] { 0x03 },
+                new byte[] { 0x04 },
+                new byte[] { 0x05 }
+            })
+        ]);
+        await binaryCollection.FlushAsync();
+
+        await binaryCollection.CreateIndexAsync("binary_vector", IndexType.BinFlat, SimilarityMetricType.Jaccard);
+        await binaryCollection.WaitForIndexBuildAsync("binary_vector");
+        await binaryCollection.LoadAsync();
+
+        var queryVector = new ReadOnlyMemory<byte>[] { new byte[] { 0x01 } };
+        var iterator = binaryCollection.SearchWithIteratorAsync(
+            "binary_vector",
+            queryVector,
+            SimilarityMetricType.Jaccard,
+            limit: 5,
+            batchSize: 2);
+
+        List<SearchResults> results = new();
+        await foreach (var result in iterator)
+        {
+            results.Add(result);
+        }
+
+        int totalResults = results.Sum(r => r.Ids.LongIds!.Count);
+        Assert.Equal(5, totalResults);
+    }
+
+    public async Task InitializeAsync()
+    {
+        await Collection.DropAsync();
+        await Client.CreateCollectionAsync(
+            CollectionName,
+            new[]
+            {
+                FieldSchema.Create<long>("id", isPrimaryKey: true),
+                FieldSchema.CreateFloatVector("float_vector", 4),
+            });
+    }
+
+    public Task DisposeAsync()
+    {
+        Client.Dispose();
+        return Task.CompletedTask;
+    }
+
+    private const string CollectionName = nameof(SearchWithIteratorExtendedTests);
+
+    private readonly MilvusClient Client;
+
+    private MilvusCollection Collection { get; }
+
+    public SearchWithIteratorExtendedTests(MilvusFixture milvusFixture)
+    {
+        Client = milvusFixture.CreateClient();
+        Collection = Client.GetCollection(CollectionName);
+    }
+}