Add support for warnings in FT.SEARCH, AGGREGATE and HYBRID

src/NRedisStack/PublicAPI/PublicAPI.Unshipped.txt

@@ -43,3 +43,7 @@ NRedisStack.TimeSeriesCommandsAsync.MRangeAsync(NRedisStack.DataTypes.TimeStamp
 NRedisStack.TimeSeriesCommandsAsync.MRevRangeAsync(NRedisStack.DataTypes.TimeStamp fromTimeStamp, NRedisStack.DataTypes.TimeStamp toTimeStamp, System.Collections.Generic.IReadOnlyCollection<string!>! filter, bool latest = false, System.Collections.Generic.IReadOnlyCollection<NRedisStack.DataTypes.TimeStamp>? filterByTs = null, (long, long)? filterByValue = null, bool? withLabels = null, System.Collections.Generic.IReadOnlyCollection<string!>? selectLabels = null, long? count = null, NRedisStack.DataTypes.TimeStamp? align = null, NRedisStack.TsAggregations aggregation = default(NRedisStack.TsAggregations), long? timeBucket = null, NRedisStack.Literals.Enums.TsBucketTimestamps? bt = null, bool empty = false, (string!, NRedisStack.Literals.Enums.TsReduce)? groupbyTuple = null) -> System.Threading.Tasks.Task<System.Collections.Generic.IReadOnlyList<(string! key, System.Collections.Generic.IReadOnlyList<NRedisStack.DataTypes.TimeSeriesLabel!>! labels, System.Collections.Generic.IReadOnlyList<NRedisStack.DataTypes.TimeSeriesTuple!>! values)>!>!
 NRedisStack.TimeSeriesCommandsAsync.RangeAsync(string! key, NRedisStack.DataTypes.TimeStamp fromTimeStamp, NRedisStack.DataTypes.TimeStamp toTimeStamp, bool latest = false, System.Collections.Generic.IReadOnlyCollection<NRedisStack.DataTypes.TimeStamp>? filterByTs = null, (long, long)? filterByValue = null, long? count = null, NRedisStack.DataTypes.TimeStamp? align = null, NRedisStack.TsAggregations aggregation = default(NRedisStack.TsAggregations), long? timeBucket = null, NRedisStack.Literals.Enums.TsBucketTimestamps? bt = null, bool empty = false) -> System.Threading.Tasks.Task<System.Collections.Generic.IReadOnlyList<NRedisStack.DataTypes.TimeSeriesTuple!>!>!
 NRedisStack.TimeSeriesCommandsAsync.RevRangeAsync(string! key, NRedisStack.DataTypes.TimeStamp fromTimeStamp, NRedisStack.DataTypes.TimeStamp toTimeStamp, bool latest = false, System.Collections.Generic.IReadOnlyCollection<NRedisStack.DataTypes.TimeStamp>? filterByTs = null, (long, long)? filterByValue = null, long? count = null, NRedisStack.DataTypes.TimeStamp? align = null, NRedisStack.TsAggregations aggregation = default(NRedisStack.TsAggregations), long? timeBucket = null, NRedisStack.Literals.Enums.TsBucketTimestamps? bt = null, bool empty = false) -> System.Threading.Tasks.Task<System.Collections.Generic.IReadOnlyList<NRedisStack.DataTypes.TimeSeriesTuple!>!>!
+NRedisStack.Search.SearchResult.Warnings.get -> string![]!
+NRedisStack.Search.AggregationResult.Warnings.get -> string![]!
+NRedisStack.Search.HybridSearchResult.Warnings.get -> string![]!
+NRedisStack.Search.AggregationRequest.AddScores() -> NRedisStack.Search.AggregationRequest!

src/NRedisStack/ResponseParser.cs

@@ -1289,6 +1289,34 @@ internal static T[] ParseSearchResultsMap<T>(RedisResult result, Func<string[],
         return results;
     }
 
+    // Extracts the top-level "warning" field from a RESP3 FT.SEARCH / FT.AGGREGATE reply map.
+    // Warnings sit alongside "results"/"total_results" as a (possibly empty) array of strings.
+    internal static string[] ParseWarnings(RedisResult root)
+    {
+        var arr = (RedisResult[])root!;
+        for (int i = 0; i + 1 < arr.Length; i += 2)
+        {
+            if ((string)arr[i]! == "warning" && arr[i + 1].Resp3Type is ResultType.Array)
+            {
+                var raw = (RedisResult[])arr[i + 1]!;
+                if (raw.Length == 0)
+                {
+                    return [];
+                }
+
+                var warnings = new string[raw.Length];
+                for (int j = 0; j < raw.Length; j++)
+                {
+                    warnings[j] = raw[j].ToString();
+                }
+
+                return warnings;
+            }
+        }
+
+        return [];
+    }
+
     internal static JsonType[] ParseJsonTypeArray(RedisResult result)
     {
         // RESP3 adds a layer of wrapping

src/NRedisStack/Search/AggregationRequest.cs

@@ -24,6 +24,17 @@ public AggregationRequest Verbatim()
         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
+    /// <c>SORTBY</c> or <c>APPLY</c>.
+    /// </summary>
+    public AggregationRequest AddScores()
+    {
+        args.Add(SearchArgs.ADDSCORES);
+        return this;
+    }
+
     public AggregationRequest Load(params FieldName[] fields)
     {
         if (fields.Length > 0)

src/NRedisStack/Search/AggregationResult.cs

@@ -27,6 +27,13 @@ internal WithCursorAggregationResult(string indexName, RedisResult result, long
 
     public long CursorId { get; }
 
+    /// <summary>
+    /// Warnings reported by the server for this query (for example when a query times out under the
+    /// <c>return</c>/<c>return-strict</c> on-timeout policy and partial results are returned).
+    /// Only populated on RESP3; on RESP2 FT.AGGREGATE does not carry warnings and this array is empty.
+    /// </summary>
+    public string[] Warnings { get; } = [];
+
     internal AggregationResult(RedisResult result, long cursorId = -1)
     {
         var arr = (RedisResult[])result!;
@@ -35,6 +42,7 @@ internal AggregationResult(RedisResult result, long cursorId = -1)
         {
             results = ResponseParser.ParseSearchResultsMap(result, ParseRecordFromMap, out long totalResults);
             // TotalResults = totalResults; // we ignore this in RESP3 mode for consistent API behaviour
+            Warnings = ResponseParser.ParseWarnings(result);
         }
         else
         {

src/NRedisStack/Search/HybridSearchResult.cs

@@ -30,7 +30,6 @@ internal static HybridSearchResult Parse(RedisResult? result)
                             case ResultKey.ExecutionTime:
                                 obj.ExecutionTime = TimeSpan.FromSeconds((double)value);
                                 break;
-                            /* // defer Warnings until we've seen examples
                             case ResultKey.Warnings when value.Length > 0:
                                 var warnings = new string[value.Length];
                                 for (int j = 0; j < value.Length; j++)
@@ -39,7 +38,6 @@ internal static HybridSearchResult Parse(RedisResult? result)
                                 }
                                 obj.Warnings = warnings;
                                 break;
-                                */
                             case ResultKey.Results when value.Length > 0:
                                 obj._rawResults = value.ToArray();
                                 break;
@@ -118,8 +116,12 @@ private enum ResultKey
     /// </summary>
     public TimeSpan ExecutionTime { get; private set; }
 
-    // not exposing this until I've seen it being used
-    internal string[] Warnings { get; private set; } = [];
+    /// <summary>
+    /// Warnings reported by the server for this query (for example a timeout warning under the
+    /// <c>return</c>/<c>return-strict</c> on-timeout policy). FT.HYBRID reports warnings on both
+    /// RESP2 and RESP3.
+    /// </summary>
+    public string[] Warnings { get; private set; } = [];
 
     private RedisResult[] _rawResults = [];
     private Document[]? _docResults;

src/NRedisStack/Search/Literals/CommandArgs.cs

@@ -2,6 +2,7 @@ namespace NRedisStack.Search.Literals;
 
 internal class SearchArgs
 {
+    public const string ADDSCORES = "ADDSCORES";
     public const string AGGREGATE = "AGGREGATE";
     public const string APPLY = "APPLY";
     public const string AS = "AS";

src/NRedisStack/Search/SearchResult.cs

@@ -13,6 +13,13 @@ public class SearchResult
     public long TotalResults { get; }
     public List<Document> Documents { get; }
 
+    /// <summary>
+    /// Warnings reported by the server for this query (for example when a query times out under the
+    /// <c>return</c>/<c>return-strict</c> on-timeout policy and partial results are returned).
+    /// Only populated on RESP3; on RESP2 FT.SEARCH does not carry warnings and this array is empty.
+    /// </summary>
+    public string[] Warnings { get; } = [];
+
     /// <summary>
     /// Converts the documents to a list of json strings. only works on a json documents index.
     /// </summary>
@@ -61,6 +68,7 @@ internal SearchResult(RedisResult root, bool hasContent, bool hasScores, bool ha
                     return Document.Load(id, score, payload, fields, scoreExplained);
                 }, out long totalResults));
             TotalResults = totalResults;
+            Warnings = ResponseParser.ParseWarnings(root);
         }
         else // RESP2
         {

tests/NRedisStack.Tests/CustomAssertions.cs

@@ -1,3 +1,4 @@
+using NRedisStack.Search;
 using Xunit;
 
 namespace NRedisStack.Tests;
@@ -17,4 +18,38 @@ public static void LessThan<T>(T actual, T expected) where T : IComparable<T>
         Assert.True(actual.CompareTo(expected) < 0,
             $"Failure: Expected value to be less than {expected}, but found {actual}.");
     }
+
+    /// <summary>
+    /// Asserts that RediSearch has indexed exactly <paramref name="expected"/> documents in <paramref name="index"/>.
+    /// Indexing can lag behind the writes, so FT.INFO num_docs is re-read a few times to let indexing catch up before the
+    /// exact-equality assertion.
+    /// </summary>
+    public static void AssertIndexSize(ISearchCommands ft, string index, long expected)
+    {
+        long indexed = -1;
+        // allow search time to catch up
+        for (int i = 0; i < 10; i++)
+        {
+            indexed = ft.Info(index).NumDocs;
+
+            if (indexed == expected)
+                break;
+        }
+        Assert.Equal(expected, indexed);
+    }
+
+    /// <inheritdoc cref="AssertIndexSize"/>
+    public static async Task AssertIndexSizeAsync(ISearchCommandsAsync ft, string index, long expected)
+    {
+        long indexed = -1;
+        // allow search time to catch up
+        for (int i = 0; i < 10; i++)
+        {
+            indexed = (await ft.InfoAsync(index)).NumDocs;
+
+            if (indexed == expected)
+                break;
+        }
+        Assert.Equal(expected, indexed);
+    }
 }

tests/NRedisStack.Tests/Search/HybridSearchIntegrationTests.cs

@@ -7,6 +7,8 @@
 using NRedisStack.Search.Aggregation;
 using StackExchange.Redis;
 using Xunit;
+using static NRedisStack.Tests.CustomAssertions;
+using static NRedisStack.Tests.Search.SearchTestUtils;
 
 namespace NRedisStack.Tests.Search;
 
@@ -432,4 +434,69 @@ static void WriteEscaped(ReadOnlySpan<byte> span, StringBuilder sb)
             sb.Append('"');
         }
     }
+
+    // Number of documents indexed by the on-timeout test. Enough that the hybrid query cannot
+    // complete within the 1ms per-query timeout, so the on-timeout policy is guaranteed to kick in.
+    private const int TimeoutDocCount = 10_000;
+
+    // FT.HYBRID counterpart of the FT.SEARCH/FT.AGGREGATE on-timeout tests: under the `return`
+    // policy a timed-out query must return (possibly partial) results plus a warning rather than
+    // failing. FT.HYBRID reports warnings on both RESP2 and RESP3, so no protocol skip is needed.
+    // See CAE-3003.
+    [SkipIfRedisTheory(Comparison.LessThan, "8.9.0", deferVersionCheck: true)]
+    [MemberData(nameof(EndpointsFixture.Env.AllEnvironments), MemberType = typeof(EndpointsFixture.Env))]
+    public async Task TestHybridOnTimeoutReturnPopulatesWarnings(string endpointId)
+    {
+        var api = await CreateIndexAsync(endpointId, populate: false);
+        if (api.IsNull) return;
+#if NET
+        // Bulk-load with fire-and-forget so many documents load quickly, then Ping so all writes
+        // have been processed before we wait for indexing.
+        var rand = new Random(12345);
+        using var vectorData = VectorData.Lease<Half>(V1DIM);
+        for (int i = 0; i < TimeoutDocCount; i++)
+        {
+            foreach (ref Half val in vectorData.Span)
+            {
+                val = (Half)rand.NextDouble();
+            }
+
+            HashEntry[] entry =
+            [
+                new("text1", $"hello world {i}"),
+                new("vector1", vectorData.AsRedisValue())
+            ];
+            api.DB.HashSet($"{api.Index}_e{i}", entry, flags: CommandFlags.FireAndForget);
+        }
+
+        api.DB.Ping();
+
+        AssertIndexSize(api.FT, api.Index, TimeoutDocCount);
+
+        // Ensure the RETURN policy is active on every primary. A stored vector guarantees the query
+        // vector matches the field type/dimension.
+        SetSearchOnTimeout(api.DB.Multiplexer, "return");
+
+        var hash = (await api.DB.HashGetAllAsync($"{api.Index}_e0")).ToDictionary(k => k.Name, v => v.Value);
+        var vec = (byte[])hash["vector1"]!;
+        var query = new HybridSearchQuery()
+            .Search("hello world")
+            .VectorSearch(new("@vector1", VectorData.Raw(vec),
+                method: VectorSearchMethod.NearestNeighbour(TimeoutDocCount)))
+            .Combine(HybridSearchQuery.Combiner.Linear(0.5, 0.5, window: TimeoutDocCount))
+            .ReturnFields("@text1")
+            .Apply(new("format(\"%s%s%s\",@text1,@text1,@text1)", "a"), new("format(\"%s%s%s\",@a,@a,@a)", "a"),
+                new("format(\"%s%s%s\",@a,@a,@a)", "a"), new("format(\"%s%s%s\",@a,@a,@a)", "a"),
+                new("format(\"%s%s%s\",@a,@a,@a)", "a"))
+            .SortBy(SortedField.Desc("@a"))
+            .Timeout(TimeSpan.FromMilliseconds(1));
+
+        var result = api.FT.HybridSearch(api.Index, query);
+        Assert.NotNull(result);
+        Assert.NotEmpty(result.Warnings);
+        Assert.Contains(result.Warnings, w => w.ToLowerInvariant().Contains("timeout"));
+#else
+        await Task.CompletedTask;
+#endif
+    }
 }

tests/NRedisStack.Tests/Search/SearchTestUtils.cs

@@ -0,0 +1,21 @@
+using StackExchange.Redis;
+
+namespace NRedisStack.Tests.Search;
+
+public static class SearchTestUtils
+{
+    // Sets the global search-on-timeout policy on every primary. On a cluster the policy must be
+    // present on all shards (and the coordinator) for the timeout behaviour to be consistent, so we
+    // broadcast rather than target a single node.
+    public static void SetSearchOnTimeout(IConnectionMultiplexer muxer, string value)
+    {
+        foreach (var endpoint in muxer.GetEndPoints())
+        {
+            var server = muxer.GetServer(endpoint);
+            if (!server.IsReplica)
+            {
+                server.ConfigSet("search-on-timeout", value);
+            }
+        }
+    }
+}