Support WITHCOUNT / WITHOUTCOUNT on FT.SEARCH and FT.AGGREGATE

src/NRedisStack/PublicAPI/PublicAPI.Unshipped.txt

@@ -1 +1,5 @@
 #nullable enable
+NRedisStack.Search.AggregationRequest.WithCount(bool? value = true) -> NRedisStack.Search.AggregationRequest!
+NRedisStack.Search.Query.SetWithCount(bool? value = true) -> NRedisStack.Search.Query!
+NRedisStack.Search.Query.WithCount.get -> bool?
+NRedisStack.Search.Query.WithCount.set -> void

src/NRedisStack/Search/AggregationRequest.cs

@@ -7,6 +7,7 @@ public class AggregationRequest : IDialectAwareParam
 {
     private readonly List<object> args = []; // Check if Readonly
     private bool isWithCursor = false;
+    private bool? withCount = null;
 
     public int? dialect { get; private set; } = null;
 
@@ -24,6 +25,23 @@ public AggregationRequest Verbatim()
         return this;
     }
 
+    /// <summary>
+    /// Whether the aggregation counts every result: <c>true</c> sends <c>WITHCOUNT</c>, which keeps the total
+    /// exact; <c>false</c> sends <c>WITHOUTCOUNT</c>, which lets an aggregation sorted by a numeric field stop
+    /// once it has enough results, at the cost of an exact total; <c>null</c> sends neither and leaves it to
+    /// the server.
+    /// </summary>
+    /// <seealso href="https://redis.io/docs/latest/commands/ft.aggregate/"/>
+    /// <remarks>
+    /// The server only accepts the option before the first pipeline step, so it is sent right after the
+    /// query wherever in the chain this is called; a later call replaces an earlier one.
+    /// </remarks>
+    public AggregationRequest WithCount(bool? value = true)
+    {
+        withCount = value;
+        return this;
+    }
+
     /// <summary>
     /// Exposes the full-text score values to the aggregation pipeline via the <c>ADDSCORES</c> option.
     /// The score of each result is then available as the <c>@__score</c> field, e.g. for use in
@@ -184,6 +202,21 @@ private void Dialect()
         }
     }
 
+    /// <summary>
+    /// Adds the request's arguments to <paramref name="target"/>, with the count option right after the query,
+    /// since the server rejects it after a pipeline step. The request's own arguments are left as they are,
+    /// so building the command again reflects the current option.
+    /// </summary>
+    internal void AddArgsTo(List<object> target)
+    {
+        target.Add(args[0]);
+        if (withCount.HasValue)
+        {
+            target.Add(withCount.Value ? SearchArgs.WITHCOUNT : SearchArgs.WITHOUTCOUNT);
+        }
+        target.AddRange(args.Skip(1));
+    }
+
     public List<object> GetArgs()
     {
         return args;

src/NRedisStack/Search/Literals/CommandArgs.cs

@@ -68,7 +68,9 @@ internal class SearchArgs
     public const string UNF = "UNF";
     public const string VERBATIM = "VERBATIM";
     public const string WEIGHT = "WEIGHT";
+    public const string WITHCOUNT = "WITHCOUNT";
     public const string WITHCURSOR = "WITHCURSOR";
+    public const string WITHOUTCOUNT = "WITHOUTCOUNT";
     public const string WITHPAYLOADS = "WITHPAYLOADS";
     public const string WITHSCORES = "WITHSCORES";
     public const string WITHSUFFIXTRIE = "WITHSUFFIXTRIE";

src/NRedisStack/Search/Query.cs

@@ -149,6 +149,14 @@ public HighlightTags(string open, string close)
     /// Set the query to return object payloads, if any were given
     /// </summary>
     public bool WithPayloads { get; set; }
+    /// <summary>
+    /// Whether the query counts every result: <c>true</c> sends <c>WITHCOUNT</c>, which returns an exact
+    /// total; <c>false</c> sends <c>WITHOUTCOUNT</c>, which lets a query sorted by a numeric field stop
+    /// once it has enough results, at the cost of an exact total; <c>null</c> (the default) sends neither
+    /// and leaves it to the server, which counts unless the query uses <c>DIALECT 4</c>.
+    /// </summary>
+    /// <seealso href="https://redis.io/docs/latest/commands/ft.search/"/>
+    public bool? WithCount { get; set; }
 
     /// <summary>
     /// Set the query language, for stemming purposes; see http://redisearch.io for documentation on languages and stemming
@@ -237,6 +245,10 @@ internal void SerializeRedisArgs(List<object> args)
         {
             args.Add(SearchArgs.WITHPAYLOADS);
         }
+        if (WithCount.HasValue)
+        {
+            args.Add(WithCount.Value ? SearchArgs.WITHCOUNT : SearchArgs.WITHOUTCOUNT);
+        }
         if (Language != null)
         {
             args.Add(SearchArgs.LANGUAGE);
@@ -491,6 +503,18 @@ public Query SetWithPayloads()
         return this;
     }
 
+    /// <summary>
+    /// Set whether the query counts every result; see <see cref="WithCount"/>.
+    /// </summary>
+    /// <seealso href="https://redis.io/docs/latest/commands/ft.search/"/>
+    /// <param name="value"><c>true</c> for <c>WITHCOUNT</c>, <c>false</c> for <c>WITHOUTCOUNT</c>, <c>null</c> for neither</param>
+    /// <returns>the query itself</returns>
+    public Query SetWithCount(bool? value = true)
+    {
+        WithCount = value;
+        return this;
+    }
+
     /// <summary>
     /// Set the query language, for stemming purposes
     /// </summary>

src/NRedisStack/Search/SearchCommandBuilder.cs

@@ -15,7 +15,7 @@ public static SerializedCommand Aggregate(string index, AggregationRequest query
     {
         List<object> args = [index];
         query.SerializeRedisArgs();
-        args.AddRange(query.GetArgs());
+        query.AddArgsTo(args);
         return new(AggregateCategory(query), FT.AGGREGATE, args);
     }
 
@@ -207,7 +207,7 @@ public static SerializedCommand ProfileAggregate(string IndexName, AggregationRe
             : new List<object> { IndexName, SearchArgs.AGGREGATE, SearchArgs.QUERY };
 
         query.SerializeRedisArgs();
-        args.AddRange(query.GetArgs());
+        query.AddArgsTo(args);
         // profiling a cursored aggregate still allocates the cursor, so it carries the same category
         return new(AggregateCategory(query), FT.PROFILE, args);
     }

tests/NRedisStack.Tests/Search/CountOptionTests.cs

@@ -0,0 +1,128 @@
+using NRedisStack.Search;
+using NRedisStack.Search.Aggregation;
+using Xunit;
+
+namespace NRedisStack.Tests.Search;
+
+/// <summary>
+/// Unit tests for the <c>WITHCOUNT</c> / <c>WITHOUTCOUNT</c> options of <see cref="Query"/> and
+/// <see cref="AggregationRequest"/> — verify where the builders put them. These tests do not require a
+/// running server.
+/// </summary>
+public class CountOptionTests
+{
+    private static string[] SearchArgs(Query query) =>
+        SearchCommandBuilder.Search("idx", query).Args.Select(arg => arg.ToString()!).ToArray();
+
+    private static object[] AggregateArgs(AggregationRequest request) =>
+        SearchCommandBuilder.Aggregate("idx", request).Args.Skip(1).ToArray();
+
+    [Fact]
+    public void QuerySendsNoCountOptionByDefault()
+    {
+        var args = SearchArgs(new Query("foo"));
+
+        Assert.DoesNotContain("WITHCOUNT", args);
+        Assert.DoesNotContain("WITHOUTCOUNT", args);
+    }
+
+    [Fact]
+    public void QueryWithCount()
+    {
+        var args = SearchArgs(new Query("foo").SetWithCount());
+
+        Assert.Contains("WITHCOUNT", args);
+        Assert.DoesNotContain("WITHOUTCOUNT", args);
+    }
+
+    [Fact]
+    public void QueryWithoutCount()
+    {
+        var args = SearchArgs(new Query("foo").SetWithCount(false));
+
+        Assert.Contains("WITHOUTCOUNT", args);
+        Assert.DoesNotContain("WITHCOUNT", args);
+    }
+
+    [Fact]
+    public void QueryCountOptionCanBeCleared()
+    {
+        var args = SearchArgs(new Query("foo").SetWithCount(false).SetWithCount(null));
+
+        Assert.DoesNotContain("WITHCOUNT", args);
+        Assert.DoesNotContain("WITHOUTCOUNT", args);
+    }
+
+    [Fact]
+    public void AggregationCountOptionFollowsQueryWhenSetLast()
+    {
+        // The server rejects the option after a pipeline step such as LOAD, so it has to follow the
+        // query even when it is set last.
+        var request = new AggregationRequest("@n:[0 10]")
+            .Load(FieldName.Of("@n"))
+            .SortBy(10, SortedField.Asc("@n"))
+            .WithCount(false);
+
+        object[] expected =
+        [
+            "@n:[0 10]", "WITHOUTCOUNT",
+            "LOAD", 1, "@n",
+            "SORTBY", 2, "@n", "ASC", "MAX", 10
+        ];
+
+        Assert.Equal(expected, AggregateArgs(request));
+    }
+
+    [Fact]
+    public void AggregationLaterCountOptionReplacesEarlierOne()
+    {
+        var request = new AggregationRequest("*")
+            .WithCount(false)
+            .SortBy(SortedField.Asc("@n"))
+            .WithCount();
+
+        object[] expected = ["*", "WITHCOUNT", "SORTBY", 2, "@n", "ASC"];
+
+        Assert.Equal(expected, AggregateArgs(request));
+    }
+
+    [Fact]
+    public void AggregationCountOptionCanBeCleared()
+    {
+        var request = new AggregationRequest("*")
+            .WithCount(false)
+            .SortBy(SortedField.Asc("@n"))
+            .WithCount(null);
+
+        object[] expected = ["*", "SORTBY", 2, "@n", "ASC"];
+
+        Assert.Equal(expected, AggregateArgs(request));
+    }
+
+    [Fact]
+    public void AggregationCountOptionChangesTakeEffectWhenRequestIsReused()
+    {
+        var request = new AggregationRequest("*").WithCount(false);
+        AggregateArgs(request);
+
+        request.WithCount(null);
+        object[] cleared = ["*"];
+        Assert.Equal(cleared, AggregateArgs(request));
+
+        request.WithCount();
+        object[] counted = ["*", "WITHCOUNT"];
+        Assert.Equal(counted, AggregateArgs(request));
+    }
+
+    [Fact]
+    public void AggregationCountOptionLeavesCallerInsertedArgsIntact()
+    {
+        var request = new AggregationRequest("*").WithCount();
+        request.GetArgs().Insert(1, "VERBATIM");
+        request.WithCount(false);
+
+        object[] expected = ["*", "WITHOUTCOUNT", "VERBATIM"];
+
+        Assert.Equal(expected, AggregateArgs(request));
+    }
+}

tests/NRedisStack.Tests/Search/SearchTests.cs

@@ -160,6 +160,72 @@ public void TestAggregationRequestVerbatim(string endpointId)
         Assert.Equal(0, res.TotalResults);
     }
 
+    [SkipIfRedisTheory(Is.Enterprise, Comparison.LessThan, "7.2.0")]
+    [MemberData(nameof(EndpointsFixture.Env.AllEnvironments), MemberType = typeof(EndpointsFixture.Env))]
+    public void TestCountOptionsSortedByNumericField(string endpointId)
+    {
+        SkipClusterPre8(endpointId);
+        IDatabase db = GetCleanDatabase(endpointId);
+        var ft = db.FT();
+        Schema sc = new Schema().AddNumericField("n", sortable: true);
+        ft.Create(index, FTCreateParams.CreateParams(), sc);
+        for (int i = 0; i < 100; i++)
+        {
+            AddDocument(db, new Document($"doc{i}").Set("n", i));
+        }
+        AssertIndexSize(ft, index, 100);
+
+        Query Page() => new Query("@n:[(50 +inf]").SetSortBy("n", ascending: true).Limit(0, 3).SetNoContent()
+            .Dialect(2);
+        var counted = ft.Search(index, Page().SetWithCount());
+        Assert.Equal(49, counted.TotalResults);
+        Assert.Equal(["doc51", "doc52", "doc53"], counted.Documents.Select(doc => doc.Id));
+        var uncounted = ft.Search(index, Page().SetWithCount(false));
+        Assert.Equal(["doc51", "doc52", "doc53"], uncounted.Documents.Select(doc => doc.Id));
+
+        // Set after SORTBY and LOAD, where the server would reject it, to show that it is sent right
+        // after the query regardless.
+        var aggregation = ft.Aggregate(index, new AggregationRequest("@n:[(50 +inf]")
+            .SortBy(3, SortedField.Asc("@n"))
+            .Load(FieldName.Of("@n"))
+            .WithCount(false)
+            .Dialect(2));
+        Assert.Equal([51L, 52L, 53L], Enumerable.Range(0, 3).Select(i => aggregation.GetRow(i).GetLong("n")));
+    }
+
+    [SkipIfRedisTheory(Is.Enterprise, Comparison.LessThan, "7.2.0")]
+    [MemberData(nameof(EndpointsFixture.Env.AllEnvironments), MemberType = typeof(EndpointsFixture.Env))]
+    public async Task TestCountOptionsSortedByNumericFieldAsync(string endpointId)
+    {
+        SkipClusterPre8(endpointId);
+        IDatabase db = GetCleanDatabase(endpointId);
+        var ft = db.FT();
+        Schema sc = new Schema().AddNumericField("n", sortable: true);
+        await ft.CreateAsync(index, FTCreateParams.CreateParams(), sc);
+        for (int i = 0; i < 100; i++)
+        {
+            AddDocument(db, new Document($"doc{i}").Set("n", i));
+        }
+        AssertIndexSize(ft, index, 100);
+
+        Query Page() => new Query("@n:[(50 +inf]").SetSortBy("n", ascending: true).Limit(0, 3).SetNoContent()
+            .Dialect(2);
+        var counted = await ft.SearchAsync(index, Page().SetWithCount());
+        Assert.Equal(49, counted.TotalResults);
+        Assert.Equal(["doc51", "doc52", "doc53"], counted.Documents.Select(doc => doc.Id));
+        var uncounted = await ft.SearchAsync(index, Page().SetWithCount(false));
+        Assert.Equal(["doc51", "doc52", "doc53"], uncounted.Documents.Select(doc => doc.Id));
+
+        // Set after SORTBY and LOAD, where the server would reject it, to show that it is sent right
+        // after the query regardless.
+        var aggregation = await ft.AggregateAsync(index, new AggregationRequest("@n:[(50 +inf]")
+            .SortBy(3, SortedField.Asc("@n"))
+            .Load(FieldName.Of("@n"))
+            .WithCount(false)
+            .Dialect(2));
+        Assert.Equal([51L, 52L, 53L], Enumerable.Range(0, 3).Select(i => aggregation.GetRow(i).GetLong("n")));
+    }
+
     [SkipIfRedisTheory(Is.Enterprise)]
     [MemberData(nameof(EndpointsFixture.Env.AllEnvironments), MemberType = typeof(EndpointsFixture.Env))]
     public async Task TestAggregationRequestVerbatimAsync(string endpointId)