feat(search): Add IP (Inner Product) distance metric support for vector similarity (#5559)

Fixed: #5556

Author: vyavdoshenko

src/core/search/base.h

@@ -21,7 +21,7 @@ namespace dfly::search {
 
 using DocId = uint32_t;
 
-enum class VectorSimilarity { L2, COSINE };
+enum class VectorSimilarity { L2, IP, COSINE };
 
 using OwnedFtVector = std::pair<std::unique_ptr<float[]>, size_t /* dimension (size) */>;
 

src/core/search/search_test.cc

@@ -814,6 +814,41 @@ TEST_P(KnnTest, Cosine) {
   }
 }
 
+TEST_P(KnnTest, IP) {
+  // Test with normalized unit vectors for IP distance
+  // Using unit vectors pointing in different directions
+  const pair<float, float> kTestCoords[] = {
+      {1.0f, 0.0f}, {0.0f, 1.0f}, {-1.0f, 0.0f}, {0.0f, -1.0f}};
+
+  auto schema = MakeSimpleSchema({{"pos", SchemaField::VECTOR}});
+  schema.fields["pos"].special_params =
+      SchemaField::VectorParams{GetParam(), 2, VectorSimilarity::IP};
+  FieldIndices indices{schema, kEmptyOptions, PMR_NS::get_default_resource(), nullptr};
+
+  for (size_t i = 0; i < ABSL_ARRAYSIZE(kTestCoords); i++) {
+    string coords = ToBytes({kTestCoords[i].first, kTestCoords[i].second});
+    MockedDocument doc{Map{{"pos", coords}}};
+    indices.Add(i, doc);
+  }
+
+  SearchAlgorithm algo{};
+  QueryParams params;
+
+  // Query with vector pointing right - should find exact match (highest dot product)
+  {
+    params["vec"] = ToBytes({1.0f, 0.0f});
+    algo.Init("* =>[KNN 1 @pos $vec]", &params);
+    EXPECT_THAT(algo.Search(&indices).ids, testing::UnorderedElementsAre(0));
+  }
+
+  // Query with vector pointing up - should find exact match (highest dot product)
+  {
+    params["vec"] = ToBytes({0.0f, 1.0f});
+    algo.Init("* =>[KNN 1 @pos $vec]", &params);
+    EXPECT_THAT(algo.Search(&indices).ids, testing::UnorderedElementsAre(1));
+  }
+}
+
 TEST_P(KnnTest, AddRemove) {
   auto schema = MakeSimpleSchema({{"pos", SchemaField::VECTOR}});
   schema.fields["pos"].special_params =
@@ -894,12 +929,22 @@ TEST_F(SearchTest, VectorDistanceBasic) {
   EXPECT_GE(cos_dist, 0.0f);
   EXPECT_LE(cos_dist, 2.0f);  // Cosine distance range
 
+  // Test IP distance
+  float ip_dist = VectorDistance(vec1.data(), vec2.data(), 3, VectorSimilarity::IP);
+  // IP distance can be negative for non-normalized vectors
+  EXPECT_NE(ip_dist, 0.0f);  // Should be non-zero for different vectors
+
   // Test identical vectors
   float l2_same = VectorDistance(vec1.data(), vec1.data(), 3, VectorSimilarity::L2);
   EXPECT_NEAR(l2_same, 0.0f, 1e-6);
 
   float cos_same = VectorDistance(vec1.data(), vec1.data(), 3, VectorSimilarity::COSINE);
   EXPECT_NEAR(cos_same, 0.0f, 1e-6);
+
+  float ip_same = VectorDistance(vec1.data(), vec1.data(), 3, VectorSimilarity::IP);
+  // For identical vectors: IP = 1 - dot_product(v, v) = 1 - ||v||^2
+  // For vec1 = {1, 2, 3}: ||v||^2 = 1 + 4 + 9 = 14, so IP = 1 - 14 = -13
+  EXPECT_LT(ip_same, 0.0f);  // Should be negative for non-normalized vectors
 }
 
 TEST_F(SearchTest, VectorDistanceConsistency) {
@@ -914,6 +959,10 @@ TEST_F(SearchTest, VectorDistanceConsistency) {
   float cos_dist1 = VectorDistance(vec1.data(), vec2.data(), 5, VectorSimilarity::COSINE);
   float cos_dist2 = VectorDistance(vec1.data(), vec2.data(), 5, VectorSimilarity::COSINE);
   EXPECT_EQ(cos_dist1, cos_dist2);
+
+  float ip_dist1 = VectorDistance(vec1.data(), vec2.data(), 5, VectorSimilarity::IP);
+  float ip_dist2 = VectorDistance(vec1.data(), vec2.data(), 5, VectorSimilarity::IP);
+  EXPECT_EQ(ip_dist1, ip_dist2);
 }
 
 static void BM_VectorSearch(benchmark::State& state) {

src/core/search/vector_utils.cc

@@ -29,7 +29,15 @@ FAST_MATH float L2Distance(const float* u, const float* v, size_t dims) {
   return sqrt(sum);
 }
 
-// TODO: Normalize vectors ahead if cosine distance is used
+// Inner product distance: 1 - dot_product(u, v)
+// For normalized vectors, this is equivalent to cosine distance
+FAST_MATH float IPDistance(const float* u, const float* v, size_t dims) {
+  float sum_uv = 0;
+  for (size_t i = 0; i < dims; i++)
+    sum_uv += u[i] * v[i];
+  return 1.0f - sum_uv;
+}
+
 FAST_MATH float CosineDistance(const float* u, const float* v, size_t dims) {
   float sum_uv = 0, sum_uu = 0, sum_vv = 0;
   for (size_t i = 0; i < dims; i++) {
@@ -71,6 +79,8 @@ float VectorDistance(const float* u, const float* v, size_t dims, VectorSimilari
   switch (sim) {
     case VectorSimilarity::L2:
       return L2Distance(u, v, dims);
+    case VectorSimilarity::IP:
+      return IPDistance(u, v, dims);
     case VectorSimilarity::COSINE:
       return CosineDistance(u, v, dims);
   };

src/server/search/doc_index.cc

@@ -148,7 +148,9 @@ string DocIndexInfo::BuildRestoreCommand() const {
     Overloaded info{
         [](monostate) {},
         [out = &out](const search::SchemaField::VectorParams& params) {
-          auto sim = params.sim == search::VectorSimilarity::L2 ? "L2" : "COSINE";
+          auto sim = params.sim == search::VectorSimilarity::L2   ? "L2"
+                     : params.sim == search::VectorSimilarity::IP ? "IP"
+                                                                  : "COSINE";
           absl::StrAppend(out, " ", params.use_hnsw ? "HNSW" : "FLAT", " 6 ", "DIM ", params.dim,
                           " DISTANCE_METRIC ", sim, " INITIAL_CAP ", params.capacity);
         },

src/server/search/search_family.cc

@@ -80,8 +80,9 @@ search::SchemaField::VectorParams ParseVectorParams(CmdArgParser* parser) {
   for (size_t i = 0; i * 2 < num_args; i++) {
     if (parser->Check("DIM", &params.dim)) {
     } else if (parser->Check("DISTANCE_METRIC")) {
-      params.sim = parser->MapNext("L2", search::VectorSimilarity::L2, "COSINE",
-                                   search::VectorSimilarity::COSINE);
+      params.sim =
+          parser->MapNext("L2", search::VectorSimilarity::L2, "IP", search::VectorSimilarity::IP,
+                          "COSINE", search::VectorSimilarity::COSINE);
     } else if (parser->Check("INITIAL_CAP", &params.capacity)) {
     } else if (parser->Check("M", &params.hnsw_m)) {
     } else if (parser->Check("EF_CONSTRUCTION", &params.hnsw_ef_construction)) {