Add filtered vector search support to LSMVectorIndex (#3072)

* Initial plan

* Add filtered search support to LSMVectorIndex with RID-based filtering

Co-authored-by: lvca <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: lvca <[email protected]>

Author: Copilot

engine/src/main/java/com/arcadedb/index/vector/LSMVectorIndex.java

@@ -140,6 +140,40 @@ public interface GraphBuildCallback {
     void onGraphBuildProgress(String phase, int processedNodes, int totalNodes, long vectorAccesses);
   }
 
+  /**
+   * Custom Bits implementation for filtering vector search by RID.
+   * Maps graph ordinals to vector IDs, then checks if the corresponding RID is in the allowed set.
+   */
+  private class RIDBitsFilter implements Bits {
+    private final Set<RID> allowedRIDs;
+    private final int[] ordinalToVectorIdSnapshot;
+    private final VectorLocationIndex vectorIndexSnapshot;
+
+    RIDBitsFilter(final Set<RID> allowedRIDs, final int[] ordinalToVectorIdSnapshot, final VectorLocationIndex vectorIndexSnapshot) {
+      this.allowedRIDs = allowedRIDs;
+      this.ordinalToVectorIdSnapshot = ordinalToVectorIdSnapshot;
+      this.vectorIndexSnapshot = vectorIndexSnapshot;
+    }
+
+    @Override
+    public boolean get(final int ordinal) {
+      // Check if ordinal is within bounds
+      if (ordinal < 0 || ordinal >= ordinalToVectorIdSnapshot.length)
+        return false;
+
+      // Map ordinal to vector ID
+      final int vectorId = ordinalToVectorIdSnapshot[ordinal];
+
+      // Get the RID for this vector ID
+      final VectorLocationIndex.VectorLocation loc = vectorIndexSnapshot.getLocation(vectorId);
+      if (loc == null || loc.deleted)
+        return false;
+
+      // Check if this RID is in the allowed set
+      return allowedRIDs.contains(loc.rid);
+    }
+  }
+
   /**
    * Comparable wrapper for float[] to use in transaction tracking.
    * Vectors are compared by their hash code for uniqueness in the transaction map.
@@ -1803,6 +1837,21 @@ private MutablePage createNewVectorDataPage(final int pageNum) {
    * @return List of pairs containing RID and similarity score
    */
   public List<Pair<RID, Float>> findNeighborsFromVector(final float[] queryVector, final int k) {
+    return findNeighborsFromVector(queryVector, k, null);
+  }
+
+  /**
+   * Search for k nearest neighbors to the given vector within a filtered set of RIDs.
+   * This method allows restricting the search space to specific records, useful for
+   * filtering by user ID, category, or other criteria during graph traversal.
+   *
+   * @param queryVector The query vector to search for
+   * @param k           The number of neighbors to return
+   * @param allowedRIDs Optional set of RIDs to restrict search to (null means no filtering)
+   *
+   * @return List of pairs containing RID and similarity score
+   */
+  public List<Pair<RID, Float>> findNeighborsFromVector(final float[] queryVector, final int k, final Set<RID> allowedRIDs) {
     if (queryVector == null)
       throw new IllegalArgumentException("Query vector cannot be null");
 
@@ -1859,14 +1908,18 @@ public List<Pair<RID, Float>> findNeighborsFromVector(final float[] queryVector,
           this  // Pass LSM index reference for quantization support
       );
 
-      // Perform search
+      // Perform search with optional RID filtering
+      final Bits bitsFilter = (allowedRIDs != null && !allowedRIDs.isEmpty())
+          ? new RIDBitsFilter(allowedRIDs, ordinalToVectorId, vectorIndex)
+          : Bits.ALL;
+
       final SearchResult searchResult = GraphSearcher.search(
           queryVectorFloat,
           k,
           vectors,
           metadata.similarityFunction,
           graphIndex,
-          Bits.ALL
+          bitsFilter
       );
 
       LogManager.instance().log(this, Level.INFO,

engine/src/test/java/com/arcadedb/index/vector/LSMVectorIndexTest.java

@@ -2101,6 +2101,113 @@ CREATE INDEX ON CycleTest (vec) LSM_VECTOR
     deleteDirectory(dbDir);
   }
 
+  @Test
+  void filteredSearchByRID() {
+    database.transaction(() -> {
+      // Create the schema
+      database.command("sql", "CREATE DOCUMENT TYPE FilteredDoc");
+      database.command("sql", "CREATE PROPERTY FilteredDoc.id STRING");
+      database.command("sql", "CREATE PROPERTY FilteredDoc.category STRING");
+      database.command("sql", "CREATE PROPERTY FilteredDoc.embedding ARRAY_OF_FLOATS");
+
+      // Create the LSM_VECTOR index
+      database.command("sql", """
+          CREATE INDEX ON FilteredDoc (embedding) LSM_VECTOR
+          METADATA {
+            "dimensions": 3,
+            "similarity": "COSINE",
+            "maxConnections": 8,
+            "beamWidth": 50
+          }""");
+    });
+
+    // Create test data with different categories
+    final List<com.arcadedb.database.RID> categoryARIDs = new ArrayList<>();
+    final List<com.arcadedb.database.RID> categoryBRIDs = new ArrayList<>();
+    
+    database.transaction(() -> {
+      for (int i = 0; i < 20; i++) {
+        final var doc = database.newDocument("FilteredDoc");
+        doc.set("id", "doc" + i);
+        doc.set("category", i < 10 ? "A" : "B");
+        
+        // Create vectors with some pattern based on category
+        final float[] vector = new float[3];
+        if (i < 10) {
+          // Category A: vectors around [1, 1, 1]
+          vector[0] = 1.0f + (i * 0.1f);
+          vector[1] = 1.0f + (i * 0.1f);
+          vector[2] = 1.0f + (i * 0.1f);
+        } else {
+          // Category B: vectors around [10, 10, 10]
+          vector[0] = 10.0f + ((i - 10) * 0.1f);
+          vector[1] = 10.0f + ((i - 10) * 0.1f);
+          vector[2] = 10.0f + ((i - 10) * 0.1f);
+        }
+        doc.set("embedding", vector);
+        
+        final com.arcadedb.database.RID rid = doc.save().getIdentity();
+        if (i < 10) {
+          categoryARIDs.add(rid);
+        } else {
+          categoryBRIDs.add(rid);
+        }
+      }
+    });
+
+    // Get the index
+    final com.arcadedb.index.TypeIndex typeIndex = (com.arcadedb.index.TypeIndex) database.getSchema()
+        .getIndexByName("FilteredDoc[embedding]");
+    final LSMVectorIndex index = (LSMVectorIndex) typeIndex.getIndexesOnBuckets()[0];
+
+    database.transaction(() -> {
+      // Query vector close to category A
+      final float[] queryVector = {1.5f, 1.5f, 1.5f};
+
+      // Test 1: Search without filter - should return results from both categories
+      final List<com.arcadedb.utility.Pair<com.arcadedb.database.RID, Float>> unfilteredResults = 
+          index.findNeighborsFromVector(queryVector, 10);
+      assertThat(unfilteredResults).as("Unfiltered search should return results").isNotEmpty();
+      assertThat(unfilteredResults.size()).as("Should return up to 10 results").isLessThanOrEqualTo(10);
+
+      // Test 2: Search with filter for category A only
+      final Set<com.arcadedb.database.RID> allowedRIDs = new HashSet<>(categoryARIDs);
+      final List<com.arcadedb.utility.Pair<com.arcadedb.database.RID, Float>> filteredResults = 
+          index.findNeighborsFromVector(queryVector, 10, allowedRIDs);
+      
+      assertThat(filteredResults).as("Filtered search should return results").isNotEmpty();
+      assertThat(filteredResults.size()).as("Should return at most 10 results").isLessThanOrEqualTo(10);
+      
+      // Verify all results are from the allowed set
+      for (final var result : filteredResults) {
+        assertThat(allowedRIDs).as("Result RID should be in allowed set").contains(result.getFirst());
+      }
+
+      // Test 3: Search with filter for category B only  
+      final Set<com.arcadedb.database.RID> categoryBSet = new HashSet<>(categoryBRIDs);
+      final List<com.arcadedb.utility.Pair<com.arcadedb.database.RID, Float>> categoryBResults = 
+          index.findNeighborsFromVector(queryVector, 10, categoryBSet);
+      
+      // Since query vector is close to category A, but we filter to category B,
+      // we should still get results (from category B), just with higher distances
+      assertThat(categoryBResults).as("Filtered search for category B should return results").isNotEmpty();
+      
+      for (final var result : categoryBResults) {
+        assertThat(categoryBSet).as("Result RID should be from category B").contains(result.getFirst());
+      }
+
+      // Test 4: Empty filter should work like unfiltered
+      final List<com.arcadedb.utility.Pair<com.arcadedb.database.RID, Float>> emptyFilterResults = 
+          index.findNeighborsFromVector(queryVector, 10, new HashSet<>());
+      assertThat(emptyFilterResults).as("Empty filter should return results like unfiltered").isNotEmpty();
+
+      // Test 5: Null filter should work like unfiltered
+      final List<com.arcadedb.utility.Pair<com.arcadedb.database.RID, Float>> nullFilterResults = 
+          index.findNeighborsFromVector(queryVector, 10, null);
+      assertThat(nullFilterResults).as("Null filter should return results like unfiltered").isNotEmpty();
+    });
+  }
+
   /**
    * Helper method to recursively delete a directory using Files.walk() API
    */