Stability wedge detection and NaN fix for vectors

integration/test_vector_range_nonfinite.py

@@ -1,8 +1,11 @@
 """VECTOR_RANGE with non-finite distances and radii.
 
-A NaN or infinite component in a stored or a query vector makes its distance
-non-finite. A NaN or +inf distance is within no radius, not even an infinite
-one, and neither is any non-finite COSINE distance, so plain, compound and
+A stored vector with a NaN or infinite component, or a magnitude too large
+for float, is rejected at ingest and left out of the vector index. A COSINE
+query vector
+like that is rejected with an error. For L2 and IP a non-finite query vector
+is accepted and makes its distances non-finite. A NaN or +inf distance is
+within no radius, not even an infinite one, so plain, compound and
 FT.AGGREGATE queries exclude the document and a negated range includes it. An
 IP distance of -inf is within every radius and reported as -inf, as on Redis,
 also through the cluster coordinator. An infinite radius matches every
@@ -11,6 +14,9 @@
 
 import struct
 
+import pytest
+from valkey.exceptions import ResponseError
+
 from utils import IndexingTestHelper
 from valkey_search_test_case import (
     ValkeySearchClusterTestCase,
@@ -26,23 +32,25 @@ def _vec(*xs):
     return struct.pack(f"<{len(xs)}f", *xs)
 
 
-# a and b are at distance 0 and 1 (2 for L2) from Q. Every distance of the
-# other documents to Q is non-finite: n holds a NaN, p and m an infinity (for
-# IP, p is at -inf and m at +inf).
+# a and b are at distance 0 and 1 (2 for L2) from Q.
 DOCS = {
     "a": _vec(1.0, 0.0, 0.0, 0.0),
     "b": _vec(0.0, 1.0, 0.0, 0.0),
+}
+# Rejected at ingest: n holds a NaN, p and m an infinity.
+REJECTED = {
     "n": _vec(NAN, 0.0, 0.0, 0.0),
     "p": _vec(INF, 0.0, 0.0, 0.0),
     "m": _vec(-INF, 0.0, 0.0, 0.0),
 }
 Q = _vec(1.0, 0.0, 0.0, 0.0)
-# Every distance to these queries is non-finite. For IP, a and p are at -inf
-# from INF_Q.
+# Every distance to these queries is non-finite. For IP, a is at -inf from
+# INF_Q.
 NAN_Q = _vec(1.0, NAN, 0.0, 0.0)
 INF_Q = _vec(INF, 0.0, 0.0, 0.0)
+NON_FINITE_QUERIES = (NAN_Q, INF_Q)
 # For IP, the documents at -inf from each query.
-IP_NEG_INF = {Q: {"p"}, NAN_Q: set(), INF_Q: {"a", "p"}}
+IP_NEG_INF = {Q: set(), NAN_Q: set(), INF_Q: {"a"}}
 YIELD = "=>{$yield_distance_as: d}"
 
 
@@ -113,21 +121,30 @@ def test_non_finite_distance(self):
             for metric in ("L2", "IP", "COSINE"):
                 index = f"{algo}_{metric}"
                 _create_index(client, index, algo, metric)
-                for key, vec in DOCS.items():
+                for key, vec in {**DOCS, **REJECTED}.items():
                     client.hset(f"{index}:{key}",
                                 mapping={"vec": vec, "tag": "x"})
+                IndexingTestHelper.wait_for_indexing_complete_on_all_nodes(
+                    [client], index)
                 for (blob, radius), finite in cases.items():
-                    neg_inf = IP_NEG_INF[blob] if metric == "IP" else set()
-                    within = finite | neg_inf
                     vr = "@vec:[VECTOR_RANGE $r $blob]"
                     params = ["PARAMS", "4", "r", radius, "blob", blob]
+                    if metric == "COSINE" and blob in NON_FINITE_QUERIES:
+                        with pytest.raises(ResponseError,
+                                           match="NaN or infinite"):
+                            _keys(client, index, vr, *params)
+                        continue
+                    neg_inf = IP_NEG_INF[blob] if metric == "IP" else set()
+                    within = finite | neg_inf
                     got = _distances(client, index, vr + YIELD, *params)
                     assert set(got) == within, (index, radius, vr, got)
                     for key in neg_inf:
                         assert got[key] == -INF, (index, radius, key, got)
+                    # A rejected vector counts as a missing field, so its key
+                    # stays in the tag index and a negated range returns it.
                     for query, expected in (
                         (vr + " @tag:{x}", within),
-                        ("-" + vr, set(DOCS) - within),
+                        ("-" + vr, (set(DOCS) | set(REJECTED)) - within),
                     ):
                         got = _keys(client, index, query, *params)
                         assert got == expected, (index, radius, query, got)
@@ -141,28 +158,34 @@ def test_ip_neg_inf_distance_through_coordinator(self):
         """Every coordinator reports an IP -inf distance as -inf, as a
         standalone server does, whichever shard holds the document. A document
         that only the tag branch of an OR returns has no VR distance and still
-        yields none."""
+        yields none.
+
+        Stored vectors must be finite, so the -inf comes from a finite query
+        whose dot product with p overflows float: 2^120 * 2^10 = 2^130. a is at
+        exactly 0 (2^120 * 2^-120 = 1), b at 1 and m at +inf."""
         cluster = self.new_cluster_client()
         coordinators = [
             self.new_client_for_primary(i) for i in range(self.CLUSTER_SIZE)
         ]
+        query = _vec(2.0**120, 0.0, 0.0, 0.0)
         vr = "@vec:[VECTOR_RANGE 0.5 $blob]" + YIELD
-        params = ["PARAMS", "2", "blob", Q]
+        params = ["PARAMS", "2", "blob", query]
         for algo in ("FLAT", "HNSW"):
             index = f"{algo}_IP"
             _create_index(cluster, index, algo, "IP")
             # One document at -inf on each shard, so that every coordinator
             # merges a local one and remote ones. b and m are outside the
-            # radius (at 1 and +inf) and match only @tag:{y}.
-            docs = {"a": ("x", DOCS["a"]), "b": ("y", DOCS["b"]),
-                    "m": ("y", DOCS["m"])}
+            # radius and match only @tag:{y}.
+            docs = {"a": ("x", _vec(2.0**-120, 0.0, 0.0, 0.0)),
+                    "b": ("y", DOCS["b"]),
+                    "m": ("y", _vec(-(2.0**10), 0.0, 0.0, 0.0))}
             shards = set()
             i = 0
             while len(shards) < self.CLUSTER_SIZE:
                 node = cluster.get_node_from_key(f"{index}:p{i}")
                 if node.name not in shards:
                     shards.add(node.name)
-                    docs[f"p{i}"] = ("x", DOCS["p"])
+                    docs[f"p{i}"] = ("x", _vec(2.0**10, 0.0, 0.0, 0.0))
                 i += 1
             for key, (tag, vec) in docs.items():
                 cluster.hset(f"{index}:{key}", mapping={"vec": vec, "tag": tag})
@@ -172,16 +195,15 @@ def test_ip_neg_inf_distance_through_coordinator(self):
             within["a"] = 0.0
             with_tag = {**within, "b": None, "m": None}
             for coordinator in coordinators:
-                for query, expected in (
+                for q, expected in (
                     (vr, within),
                     (vr + " @tag:{x}", within),
                     ("@tag:{y} | " + vr, with_tag),
                 ):
-                    got = _distances(coordinator, index, query, *params)
-                    assert got == expected, (index, query, got)
-                    got = _aggregate_distances(
-                        coordinator, index, query, *params)
-                    assert got == expected, (index, "FT.AGGREGATE", query, got)
+                    got = _distances(coordinator, index, q, *params)
+                    assert got == expected, (index, q, got)
+                    got = _aggregate_distances(coordinator, index, q, *params)
+                    assert got == expected, (index, "FT.AGGREGATE", q, got)
                 # The -inf documents sort before a, not after it as documents
                 # without a VR distance would.
                 result = coordinator.execute_command(

src/commands/ft_hybrid.cc

@@ -22,6 +22,7 @@
 #include "src/coordinator/client_pool.h"
 #include "src/expr/expr.h"
 #include "src/expr/value.h"
+#include "src/indexes/scoring/scorer.h"
 #include "src/metrics.h"
 #include "src/query/content_resolution.h"
 #include "src/query/fanout.h"
@@ -505,18 +506,25 @@ void RevalidateArmsBeforeFusion(MultiSearchParameters &params) {
     // Dropping preserves order, so only a fresh score can have put the arm out
     // of order.
     if (rescored) {
-      std::stable_sort(neighbors.begin(), neighbors.end(),
-                       [arm_score_is_distance](const indexes::Neighbor &a,
-                                               const indexes::Neighbor &b) {
-                         if (arm_score_is_distance) {
-                           if (a.distance != b.distance) {
-                             return a.distance < b.distance;
-                           }
-                         } else if (a.score != b.score) {
-                           return a.score > b.score;
-                         }
-                         return a.external_id->Str() < b.external_id->Str();
-                       });
+      // NaN sorts last, as in rank_fusion.cc: a plain `<` is not a strict weak
+      // ordering once an L2/IP query vector holding a NaN yields NaN
+      // distances, and std::stable_sort then runs outside the range.
+      std::stable_sort(
+          neighbors.begin(), neighbors.end(),
+          [arm_score_is_distance](const indexes::Neighbor &a,
+                                  const indexes::Neighbor &b) {
+            const float x = arm_score_is_distance ? a.distance : a.score;
+            const float y = arm_score_is_distance ? b.distance : b.score;
+            const bool x_nan = indexes::scoring::IsNaN(x);
+            const bool y_nan = indexes::scoring::IsNaN(y);
+            if (x_nan != y_nan) {
+              return y_nan;
+            }
+            if (!x_nan && x != y) {
+              return arm_score_is_distance ? x < y : x > y;
+            }
+            return a.external_id->Str() < b.external_id->Str();
+          });
     }
   }
 }

src/commands/ft_hybrid_parser.cc

@@ -1156,6 +1156,15 @@ absl::Status ParseFtHybridCommand(MultiSearchParameters &env,
             arm->query.size(), ") does not match index's expected size (",
             vector_index->GetVectorDataSize(), ")."));
       }
+      // Only COSINE uses the magnitude, so other metrics pay nothing here.
+      if (vector_index->GetNormalize() &&
+          !indexes::IsValidReciprocalMagnitude(indexes::CalcReciprocalMagnitude(
+              arm->query, vector_index->GetVectorDataType()))) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("Error parsing vector similarity parameters: query "
+                         "vector ",
+                         indexes::kInvalidQueryVectorError));
+      }
       // Default score_as if the user didn't YIELD_SCORE_AS.
       if (!arm->score_as) {
         auto schema_default =

src/indexes/vector_base.cc

@@ -10,6 +10,7 @@
 #include <sys/types.h>
 
 #include <algorithm>
+#include <bit>
 #include <cstddef>
 #include <cstdint>
 #include <cstdlib>
@@ -45,6 +46,7 @@
 #include "third_party/hnswlib/hnswlib.h"
 #include "vmsdk/src/log.h"
 #include "vmsdk/src/managed_pointers.h"
+#include "vmsdk/src/module_config.h"
 #include "vmsdk/src/status/status_macros.h"
 #include "vmsdk/src/utils.h"
 #include "vmsdk/src/valkey_module_api/valkey_module.h"
@@ -62,6 +64,14 @@ float CalcReciprocalMagnitude(const T *src, size_t size) {
     float v = static_cast<float>(src[i]);
     sum_sq += v * v;
   }
+  // Squares are never negative, so the sum is NaN or Inf exactly when a
+  // component is NaN or Inf or the sum overflows float. One test on the sum
+  // therefore validates the whole vector at no per-element cost. It reads the
+  // IEEE bits because -ffast-math folds std::isnan/std::isfinite to constants.
+  if (ABSL_PREDICT_FALSE((std::bit_cast<uint32_t>(sum_sq) &
+                          kFloatExponentMask) == kFloatExponentMask)) {
+    return kInvalidReciprocalMagnitude;
+  }
   return (sum_sq == 0.0f) ? 1.0f : (1.0f / std::sqrt(sum_sq));
 }
 
@@ -92,7 +102,9 @@ float CalcReciprocalMagnitude(absl::string_view record,
 template <typename T>
 std::vector<char> NormalizeVector(absl::string_view record,
                                   float reciprocal_magnitude) {
-  if (ABSL_PREDICT_FALSE(reciprocal_magnitude == 0.0f)) {
+  // <= also catches kInvalidReciprocalMagnitude, so a rejected vector is
+  // copied unscaled instead of being sign-flipped.
+  if (ABSL_PREDICT_FALSE(reciprocal_magnitude <= 0.0f)) {
     reciprocal_magnitude = 1.0f;
   }
   size_t dimensions = record.size() / sizeof(T);
@@ -232,6 +244,11 @@ absl::StatusOr<RecordResult> VectorBase::AddRecord(const InternedStringPtr &key,
     return RecordResult::kInvalidData;
   }
   auto vector_record = data.ConsumeVector();
+  // The magnitude was computed when the record was built, so rejecting a
+  // NaN/Inf or overflowing vector here costs no extra pass.
+  if (!IsValidReciprocalMagnitude(vector_record.GetReciprocalMagnitude())) {
+    return RecordResult::kInvalidData;
+  }
   float magnitude = 1.0f / vector_record.GetReciprocalMagnitude();
   VMSDK_ASSIGN_OR_RETURN(auto internal_id, TrackKey(key, magnitude));
   absl::Status add_result =
@@ -279,6 +296,11 @@ absl::StatusOr<RecordResult> VectorBase::ModifyRecord(
     return RecordResult::kInvalidData;
   }
   auto vector_record = data.ConsumeVector();
+  if (!IsValidReciprocalMagnitude(vector_record.GetReciprocalMagnitude())) {
+    [[maybe_unused]] auto res =
+        RemoveRecord(key, indexes::DeletionType::kRecord);
+    return RecordResult::kInvalidData;
+  }
   float magnitude = 1.0f / vector_record.GetReciprocalMagnitude();
   VMSDK_ASSIGN_OR_RETURN(auto internal_id, GetInternalId(key));
   VMSDK_ASSIGN_OR_RETURN(bool res,
@@ -475,7 +497,38 @@ absl::Status VectorBase::SaveTrackedKeys(
 absl::Status VectorBase::LoadTrackedKeys(
     ValkeyModuleCtx *ctx, const AttributeDataType *attribute_data_type,
     SupplementalContentChunkIter &&iter) {
+  // Keys whose stored vector is NaN/Inf or overflows, persisted before such
+  // vectors were rejected at ingest.
+  std::vector<uint64_t> rejected_ids;
+  VMSDK_RETURN_IF_ERROR(LoadTrackedKeysLocked(ctx, attribute_data_type,
+                                              std::move(iter), rejected_ids));
+  if (rejected_ids.empty()) {
+    return absl::OkStatus();
+  }
+  // Removed after key_to_metadata_mutex_ is released: RemoveRecordImpl takes
+  // resize_mutex_, which is never acquired under key_to_metadata_mutex_.
+  for (uint64_t id : rejected_ids) {
+    auto status = RemoveRecordImpl(id);
+    if (!status.ok()) {
+      VMSDK_LOG_EVERY_N_SEC(WARNING, nullptr, 1)
+          << "Failed to remove a non-finite vector during load: " << status;
+    }
+  }
+  VMSDK_LOG(WARNING, nullptr)
+      << "Left " << rejected_ids.size() << " key(s) out of vector index "
+      << vmsdk::config::RedactIfNeeded(attribute_identifier_)
+      << " during load because their vector has a NaN or infinite value or "
+         "its magnitude is too large. A later valid write re-indexes them.";
+  return absl::OkStatus();
+}
+
+absl::Status VectorBase::LoadTrackedKeysLocked(
+    ValkeyModuleCtx *ctx, const AttributeDataType *attribute_data_type,
+    SupplementalContentChunkIter &&iter, std::vector<uint64_t> &rejected_ids) {
   absl::WriterMutexLock lock(&key_to_metadata_mutex_);
+  // Placed in a rejected key's slot, so traversal through it stays finite
+  // until the slot is removed.
+  const std::vector<char> zeros(GetVectorDataSize(), 0);
 
   while (iter.HasNext()) {
     VMSDK_ASSIGN_OR_RETURN(auto metadata_str, iter.Next(),
@@ -485,12 +538,6 @@ absl::Status VectorBase::LoadTrackedKeys(
       return absl::InvalidArgumentError("Error parsing metadata from proto");
     }
     auto interned_key = StringInternStore::Intern(tracked_key_metadata.key());
-    tracked_metadata_by_key_.insert(
-        {interned_key,
-         {.internal_id = tracked_key_metadata.internal_id(),
-          .magnitude = tracked_key_metadata.magnitude()}});
-    key_by_internal_id_.insert(
-        {tracked_key_metadata.internal_id(), interned_key});
 
     auto key = vmsdk::MakeUniqueValkeyString(interned_key->Str());
     auto key_obj = vmsdk::MakeUniqueValkeyOpenKey(
@@ -508,8 +555,24 @@ absl::Status VectorBase::LoadTrackedKeys(
         interned_key, attribute_val.get(), attribute_data_type->ToProto(),
         db_num_, this);
     auto &save_vector = GetVectorLockFree(tracked_key_metadata.internal_id());
+    if (vector_record_with_size.vector_record &&
+        !IsValidReciprocalMagnitude(
+            vector_record_with_size.vector_record.GetReciprocalMagnitude())) {
+      save_vector =
+          VectorRecord::Construct(absl::string_view(zeros.data(), zeros.size()),
+                                  kDefaultMagnitude, GetVectorAllocator());
+      rejected_ids.push_back(tracked_key_metadata.internal_id());
+      continue;
+    }
+    tracked_metadata_by_key_.insert(
+        {interned_key,
+         {.internal_id = tracked_key_metadata.internal_id(),
+          .magnitude = tracked_key_metadata.magnitude()}});
+    key_by_internal_id_.insert(
+        {tracked_key_metadata.internal_id(), interned_key});
     save_vector = std::move(vector_record_with_size.vector_record);
   }
+  // Includes rejected labels, so they are never reused before removal.
   inc_id_ = GetMaxLoadedLabel() + 1;
   return absl::OkStatus();
 }
@@ -558,8 +621,15 @@ absl::StatusOr<float> VectorBase::RecomputeDistance(
   // Built with the default allocator rather than the index's own: this runs on
   // the main thread while writers may be using that allocator, and one record
   // per mutated key is not worth sharing.
-  auto vector_record = VectorRecord::Construct(
-      record, ComputeReciprocalMagnitude(record), nullptr);
+  const float reciprocal_magnitude = ComputeReciprocalMagnitude(record);
+  // A vector the index would reject at ingest has no meaningful distance.
+  if (!IsValidReciprocalMagnitude(reciprocal_magnitude)) {
+    return absl::InvalidArgumentError(
+        "Vector record contains a NaN or infinite value, or its magnitude is "
+        "too large");
+  }
+  auto vector_record =
+      VectorRecord::Construct(record, reciprocal_magnitude, nullptr);
   if (!vector_record) {
     return absl::InternalError("Could not construct a vector record");
   }