Merge remote-tracking branch 'origin/master' into feature-coordinator-hybrid-jk_update

kei-nan · kei-nan · commit 43d83c024483 · 2025-10-03T15:08:47.000+03:00
diff --git a/src/hybrid/hybrid_request.c b/src/hybrid/hybrid_request.c
@@ -107,6 +107,7 @@ int HybridRequest_BuildMergePipeline(HybridRequest *req, HybridPipelineParams *p
     HybridLookupContext *lookupCtx = InitializeHybridLookupContext(req->requests, lookup);
 
     const char *scoreAlias = params->aggregationParams.common.scoreAlias;
+    const RLookupKey *docKey = RLookup_GetKey_Read(lookup, UNDERSCORE_KEY, RLOOKUP_F_HIDDEN);
     const RLookupKey *scoreKey = NULL;
     if (scoreAlias) {
       scoreKey = RLookup_GetKey_Write(lookup, scoreAlias, RLOOKUP_F_NOFLAGS);
@@ -118,7 +119,7 @@ int HybridRequest_BuildMergePipeline(HybridRequest *req, HybridPipelineParams *p
     } else {
       scoreKey = RLookup_GetKey_Read(lookup, UNDERSCORE_SCORE, RLOOKUP_F_NOFLAGS);
     }
-    ResultProcessor *merger = RPHybridMerger_New(params->scoringCtx, depleters, req->nrequests, scoreKey, req->subqueriesReturnCodes, lookupCtx);
+    ResultProcessor *merger = RPHybridMerger_New(params->scoringCtx, depleters, req->nrequests, docKey, scoreKey, req->subqueriesReturnCodes, lookupCtx);
     params->scoringCtx = NULL; // ownership transferred to merger
     QITR_PushRP(&req->tailPipeline->qctx, merger);
     // Build the aggregation part of the tail pipeline for final result processing
diff --git a/src/hybrid/parse_hybrid.c b/src/hybrid/parse_hybrid.c
@@ -655,6 +655,10 @@ int parseHybridCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc,
   // We need a load step, implicit or an explicit one
   PLN_LoadStep *loadStep = (PLN_LoadStep *)AGPLN_FindStep(parsedCmdCtx->tailPlan, NULL, NULL, PLN_T_LOAD);
   if (!loadStep) {
+    // TBH don't think we need this implicit step, added to somehow affect the resulting response format
+    // We wanted that by default the key and score would be returned to the user
+    // This should probably be done in the hybrid send chunk where we decide on the response format.
+    // For now keeping it as is - due to time constraints
     loadStep = createImplicitLoadStep();
   } else {
     AGPLN_PopStep(&loadStep->base);
diff --git a/src/result_processor.c b/src/result_processor.c
@@ -1676,13 +1676,15 @@ dictType dictTypeHybridSearchResult = {
  typedef struct {
  ResultProcessor base;
  HybridScoringContext *hybridScoringCtx;  // Store by pointer - RPHybridMerger is responsible for freeing it
- ResultProcessor **upstreams;  // Dynamic array of upstream processors
- size_t numUpstreams;         // Number of upstream processors
- dict *hybridResults;  // keyPtr -> HybridSearchResult mapping
- dictIterator *iterator; // Iterator for yielding results
- const RLookupKey *scoreKey;  // Key for writing score as field when QEXEC_F_SEND_SCORES_AS_FIELD is set
- RPStatus* upstreamReturnCodes;  // Final return codes from each upstream
+ ResultProcessor **upstreams;     // Dynamic array of upstream processors
+ size_t numUpstreams;             // Number of upstream processors
+ dict *hybridResults;             // keyPtr -> HybridSearchResult mapping
+ dictIterator *iterator;          // Iterator for yielding results
+ const RLookupKey *scoreKey;      // Key for writing score as field when YIELD_SCORE_AS is specified
+ const RLookupKey *docKey;        // Key for reading document key when dmd is not available
+ RPStatus* upstreamReturnCodes;   // Final return codes from each upstream
  HybridLookupContext *lookupCtx;  // Lookup context for field merging
+
 } RPHybridMerger;
 
 /* Generic helper function to check if any upstream has a specific return code */
@@ -1712,24 +1714,38 @@ static inline bool RPHybridMerger_Error(const RPHybridMerger *self) {
   * @param numUpstreams - the number of upstreams
   * @param score - used to override the result's score
  */
- static void StoreUpstreamResult(SearchResult *r, dict *hybridResults, int upstreamIndex, size_t numUpstreams, double score) {
-   const char *keyPtr = SearchResult_GetDocumentMetadata(r)->keyPtr;
+ static bool hybridMergerStoreUpstreamResult(SearchResult *r, dict *hybridResults, const RLookupKey *docKey, int upstreamIndex, size_t numUpstreams, double score) {
+  // Single shard case - use dmd->keyPtr
+  const RSDocumentMetadata *dmd = SearchResult_GetDocumentMetadata(r);
+  const char *keyPtr = dmd ? dmd->keyPtr : NULL;
+  // Coordinator case - no dmd - use docKey in rlookup
+  const bool fallbackToLookup = !keyPtr && docKey;
+  if (fallbackToLookup) {
+    RSValue *docKeyValue = RLookup_GetItem(docKey, &r->rowdata);
+    if (docKeyValue != NULL) {
+      keyPtr = RSValue_StringPtrLen(docKeyValue, NULL);
+    }
+  }
+  if (!keyPtr) {
+    return false;
+  }
 
-   // Check if we've seen this document before
-   HybridSearchResult *hybridResult = (HybridSearchResult*)dictFetchValue(hybridResults, keyPtr);
+  // Check if we've seen this document before
+  HybridSearchResult *hybridResult = (HybridSearchResult*)dictFetchValue(hybridResults, keyPtr);
 
-   if (!hybridResult) {
-     // First time seeing this document - create new hybrid result
-     hybridResult = HybridSearchResult_New(numUpstreams);
-     dictAdd(hybridResults, (void*)keyPtr, hybridResult);
-   }
+  if (!hybridResult) {
+    // First time seeing this document - create new hybrid result
+    hybridResult = HybridSearchResult_New(numUpstreams);
+    dictAdd(hybridResults, (void*)keyPtr, hybridResult);
+  }
 
    SearchResult_SetScore(r, score);
    HybridSearchResult_StoreResult(hybridResult, r, upstreamIndex);
+   return true;
  }
 
  /* Helper function to consume results from a single upstream */
- static int ConsumeFromUpstream(RPHybridMerger *self, size_t maxResults, ResultProcessor *upstream, int upstreamIndex) {
+ static int hybridMergerConsumeFromUpstream(RPHybridMerger *self, size_t maxResults, ResultProcessor *upstream, int upstreamIndex) {
    size_t consumed = 0;
    int rc = RS_RESULT_OK;
    SearchResult *r = rm_calloc(1, sizeof(*r));
@@ -1739,8 +1755,12 @@ static inline bool RPHybridMerger_Error(const RPHybridMerger *self) {
        if (self->hybridScoringCtx->scoringType == HYBRID_SCORING_RRF) {
          score = consumed;
        }
-       StoreUpstreamResult(r, self->hybridResults, upstreamIndex, self->numUpstreams, score);
-       r = rm_calloc(1, sizeof(*r));
+       if (hybridMergerStoreUpstreamResult(r, self->hybridResults, self->docKey, upstreamIndex, self->numUpstreams, score)) {
+         r = rm_calloc(1, sizeof(*r));
+       } else {
+         SearchResult_Clear(r);
+         --consumed; // avoid wrong rank in RRF
+       }
    }
    RLookup_AddKeysFrom(self->lookupCtx->sourceLookups[upstreamIndex], self->lookupCtx->tailLookup, RLOOKUP_F_NOFLAGS);
    rm_free(r);
@@ -1800,7 +1820,7 @@ static inline bool RPHybridMerger_Error(const RPHybridMerger *self) {
       } else {
         window = self->hybridScoringCtx->linearCtx.window;
       }
-      int rc = ConsumeFromUpstream(self, window, self->upstreams[i], i);
+      int rc = hybridMergerConsumeFromUpstream(self, window, self->upstreams[i], i);
 
       if (rc == RS_RESULT_DEPLETING) {
         // Upstream is still active but not ready to provide results. Skip to the next.
@@ -1880,6 +1900,7 @@ static inline bool RPHybridMerger_Error(const RPHybridMerger *self) {
 ResultProcessor *RPHybridMerger_New(HybridScoringContext *hybridScoringCtx,
                                     ResultProcessor **upstreams,
                                     size_t numUpstreams,
+                                    const RLookupKey *docKey,
                                     const RLookupKey *scoreKey,
                                     RPStatus *subqueriesReturnCodes,
                                     HybridLookupContext *lookupCtx) {
@@ -1892,7 +1913,9 @@ ResultProcessor *RPHybridMerger_New(HybridScoringContext *hybridScoringCtx,
   RS_ASSERT(hybridScoringCtx);
   ret->hybridScoringCtx = hybridScoringCtx;
 
-  // Store the scoreKey for writing scores as fields when QEXEC_F_SEND_SCORES_AS_FIELD is set
+  // Store the scoreKey for reading document keys from rlookup
+  ret->docKey = docKey;
+  // Store the scoreKey for writing scores as fields when YIELD_SCORE_AS is specified
   ret->scoreKey = scoreKey;
 
   // Store reference to the hybrid request's subqueries return codes array
diff --git a/src/result_processor.h b/src/result_processor.h
@@ -307,6 +307,7 @@ StrongRef DepleterSync_New(unsigned int num_depleters, bool take_index_lock);
 ResultProcessor *RPHybridMerger_New(HybridScoringContext *hybridScoringCtx,
                                     ResultProcessor **upstreams,
                                     size_t numUpstreams,
+                                    const RLookupKey *docKey,
                                     const RLookupKey *scoreKey,
                                     RPStatus *subqueriesReturnCodes,
                                     HybridLookupContext *lookupCtx);
diff --git a/src/util/arg_parser.c b/src/util/arg_parser.c
@@ -466,7 +466,7 @@ ArgParseResult ArgParser_Parse(ArgParser *parser) {
             // Find the next unparsed positional argument
             ArgDefinition *pos_def = NULL;
             for (uint16_t pos = 1; pos <= MAX_POSITIONAL_ARGS; pos++) { // reasonable limit
-                ArgDefinition *candidate = find_positional_definition(parser, pos, NULL);
+                ArgDefinition *candidate = find_positional_definition(parser, pos, arg_name);
                 if (!candidate) break;
 
                 if (!candidate->parsed) {
diff --git a/tests/cpptests/test_cpp_hybridmerger.cpp b/tests/cpptests/test_cpp_hybridmerger.cpp
@@ -190,7 +190,7 @@ ResultProcessor* CreateLinearHybridMerger(ResultProcessor **upstreams, size_t nu
   // Create dummy return codes array for tests that don't need to track return codes
   static RPStatus dummyReturnCodes[8] = {RS_RESULT_OK}; // Static array, supports up to 8 upstreams for tests
 
-  return RPHybridMerger_New(hybridScoringCtx, upstreams, numUpstreams, NULL, dummyReturnCodes, lookupCtx);
+  return RPHybridMerger_New(hybridScoringCtx, upstreams, numUpstreams, NULL, NULL, dummyReturnCodes, lookupCtx);
 }
 
 // Helper function to create hybrid merger with RRF scoring
@@ -201,7 +201,7 @@ ResultProcessor* CreateRRFHybridMerger(ResultProcessor **upstreams, size_t numUp
   // Create dummy return codes array for tests that don't need to track return codes
   static RPStatus dummyReturnCodes[8] = {RS_RESULT_OK}; // Static array, supports up to 8 upstreams for tests
 
-  return RPHybridMerger_New(hybridScoringCtx, upstreams, numUpstreams, NULL, dummyReturnCodes, lookupCtx);
+  return RPHybridMerger_New(hybridScoringCtx, upstreams, numUpstreams, NULL, NULL, dummyReturnCodes, lookupCtx);
 }
 
 
@@ -1424,7 +1424,7 @@ TEST_F(HybridMergerTest, testUpstreamReturnCodes) {
   // Create dummy lookup context
   HybridLookupContext *lookupCtx = CreateDummyLookupContext(3);
 
-  ResultProcessor *hybridMerger = RPHybridMerger_New(hybridScoringCtx, upstreams, 3, NULL, returnCodes, lookupCtx);
+  ResultProcessor *hybridMerger = RPHybridMerger_New(hybridScoringCtx, upstreams, 3, NULL, NULL, returnCodes, lookupCtx);
 
   // Process results - this should capture the return codes
   SearchResult r = {0};
diff --git a/tests/pytests/test_hybrid_filter.py b/tests/pytests/test_hybrid_filter.py
@@ -61,7 +61,6 @@ def test_hybrid_filter_behavior():
         'FILTER', '@category:{"fruit"}'
     )
     results, _ = get_results_from_hybrid_response(response)
-    # This should return all with fruit from vector subquery (doc:1, doc:2) and all with green text (doc:3)
     env.assertEqual(set(results.keys()), {"doc:1", "doc:2", "doc:3"})
 
     response = env.cmd(
@@ -71,7 +70,6 @@ def test_hybrid_filter_behavior():
         'FILTER', '@category:{"fruit"}', "COMBINE", "RRF", "2", "CONSTANT", "30",
     )
     results, _ = get_results_from_hybrid_response(response)
-    # This should filter as before, just an extra combine
     env.assertEqual(set(results.keys()), {"doc:1", "doc:2", "doc:3"})
 
     response = env.cmd(
@@ -81,7 +79,6 @@ def test_hybrid_filter_behavior():
         "COMBINE", "RRF", "2", "CONSTANT", "30", "LOAD", 2, "__key", "category", "FILTER", "@category==\"fruit\"",
     )
     results, _ = get_results_from_hybrid_response(response)
-    # This should filter as post processing.
     env.assertEqual(set(results.keys()), {"doc:1", "doc:2"})
 
     response = env.cmd(
@@ -91,7 +88,6 @@ def test_hybrid_filter_behavior():
         'FILTER', '@category:{"vegetable"}', "COMBINE", "RRF", "2", "CONSTANT", "30", "LOAD", 2, "__key", "category", "FILTER", "@category==\"fruit\"",
     )
     results, _ = get_results_from_hybrid_response(response)
-    # This should filter as before, just an extra combine
     env.assertEqual(results, {})
 
     response = env.cmd(
@@ -101,16 +97,31 @@ def test_hybrid_filter_behavior():
         'FILTER', '@category:{"vegetable"}', "LOAD", 2, "__key", "category", "FILTER", "@category==\"clothing\"",
     )
     results, _ = get_results_from_hybrid_response(response)
-    # This should filter as before, just an extra combine
     env.assertEqual(set(results.keys()), {"doc:3"})
 
-    # post-query FILTER immediately after VSIM FILTER
+    response = env.cmd(
+        'FT.HYBRID', 'filter_idx',
+        'SEARCH', '@text:(green)',
+        'VSIM', '@vector', query_vector,
+        'FILTER', '@category:{"vegetable"}',
+    )
+    results, _ = get_results_from_hybrid_response(response)
+    env.assertEqual(set(results.keys()), {"doc:3", "doc:4"})
+
+    response = env.cmd(
+        'FT.HYBRID', 'filter_idx',
+        'SEARCH', '@text:(green)',
+        'VSIM', '@vector', query_vector,
+        'FILTER', '@category:{"vegetable"}', "FILTER", "@__key!=\"doc:3\"",
+    )
+    results, _ = get_results_from_hybrid_response(response)
+    env.assertEqual(set(results.keys()), {"doc:4"})
+
     response = env.cmd(
         'FT.HYBRID', 'filter_idx',
         'SEARCH', '@text:(green)',
         'VSIM', '@vector', query_vector,
         'FILTER', '@category:{"vegetable"}', "FILTER", "@__key==\"doc:3\"",
     )
     results, _ = get_results_from_hybrid_response(response)
-    # This should filter as before, just an extra combine
     env.assertEqual(set(results.keys()), {"doc:3"})
