Move code into reusable function

spotify · Aug 30, 2024 · a3d04c8 · a3d04c8
1 parent 12fd327
commit a3d04c8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 26 deletions.
diff --git a/cpp/src/TypedIndex.h b/cpp/src/TypedIndex.h
@@ -293,19 +293,7 @@ class TypedIndex : public Index {
   std::vector<hnswlib::labeltype>
   addItems(const std::vector<std::vector<float>> vectors,
            std::vector<hnswlib::labeltype> ids = {}, int numThreads = -1) {
-    // Convert the 2D array of float to NDArray<float, 2>
-    int numVectors = vectors.size();
-    int dimensions = numVectors > 0 ? vectors[0].size() : 0;
-    std::array<int, 2> shape = {numVectors, dimensions};
-
-    // flatten the 2d array of floats
-    std::vector<float> flatArray;
-    for (const auto &vector : vectors) {
-      flatArray.insert(flatArray.end(), vector.begin(), vector.end());
-    }
-    NDArray<float, 2> ndarray(flatArray, shape);
-
-    return addItems(ndarray, ids, numThreads);
+    return addItems(vectorsToNDArray(vectors), ids, numThreads);
   }
 
   std::vector<hnswlib::labeltype>
@@ -523,19 +511,7 @@ class TypedIndex : public Index {
   std::tuple<NDArray<hnswlib::labeltype, 2>, NDArray<dist_t, 2>>
   query(std::vector<std::vector<float>> floatQueryVectors, int k = 1,
         int numThreads = -1, long queryEf = -1) {
-    // Convert the 2D array of float to NDArray<float, 2>
-    int numVectors = floatQueryVectors.size();
-    int dimensions = numVectors > 0 ? floatQueryVectors[0].size() : 0;
-    std::array<int, 2> shape = {numVectors, dimensions};
-
-    // flatten the 2d array of floats
-    std::vector<float> flatArray;
-    for (const auto &vector : floatQueryVectors) {
-      flatArray.insert(flatArray.end(), vector.begin(), vector.end());
-    }
-    NDArray<float, 2> ndarray(flatArray, shape);
-
-    return query(ndarray, k, numThreads, queryEf);
+    return query(vectorsToNDArray(floatQueryVectors), k, numThreads, queryEf);
   }
 
   std::tuple<NDArray<hnswlib::labeltype, 2>, NDArray<dist_t, 2>>

diff --git a/cpp/src/array_utils.h b/cpp/src/array_utils.h
@@ -309,3 +309,20 @@ std::string toFloatVectorString(std::vector<data_t> vec) {
   return toFloatVectorString<dist_t, data_t, scalefactor>(vec.data(),
                                                           vec.size());
 }
+
+/**
+ * Convert a 2D vector of float to NDArray<float, 2>
+ */
+NDArray<float, 2> vectorsToNDArray(std::vector<std::vector<float>> vectors) {
+  int numVectors = vectors.size();
+  int dimensions = numVectors > 0 ? vectors[0].size() : 0;
+  std::array<int, 2> shape = {numVectors, dimensions};
+
+  // flatten the 2d array into the NDArray's underlying 1D vector
+  std::vector<float> flatArray;
+  for (const auto &vector : vectors) {
+    flatArray.insert(flatArray.end(), vector.begin(), vector.end());
+  }
+
+  return NDArray<float, 2>(flatArray, shape);
+}