From 52a17e086ac9409005d359639265da5effb53ba1 Mon Sep 17 00:00:00 2001
From: Jorge Lapa <2780099+heyjorgedev@users.noreply.github.com>
Date: Fri, 28 Feb 2025 18:21:39 +0700
Subject: [PATCH] feat: range vectors using an id prefix

---
 src/Iterators/VectorRangeIterator.php       |  1 +
 src/VectorRange.php                         | 10 ++-
 tests/Dense/Operations/RangeVectorsTest.php | 72 +++++++++++++++++++++
 3 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/src/Iterators/VectorRangeIterator.php b/src/Iterators/VectorRangeIterator.php
index e9d6d55..fef67be 100644
--- a/src/Iterators/VectorRangeIterator.php
+++ b/src/Iterators/VectorRangeIterator.php
@@ -76,6 +76,7 @@ private function fetchWithCursor(string $cursor): VectorRangeResult
         return $this->operation->range(new VectorRange(
             limit: $this->range->limit,
             cursor: $cursor,
+            prefix: $this->range->prefix,
             includeMetadata: $this->range->includeMetadata,
             includeVectors: $this->range->includeVectors,
             includeData: $this->range->includeData,
diff --git a/src/VectorRange.php b/src/VectorRange.php
index ec371a2..6cd71cc 100644
--- a/src/VectorRange.php
+++ b/src/VectorRange.php
@@ -9,6 +9,7 @@
     public function __construct(
         public int $limit,
         public string $cursor = '0',
+        public ?string $prefix = null,
         public bool $includeMetadata = false,
         public bool $includeVectors = false,
         public bool $includeData = false,
@@ -18,6 +19,7 @@ public function __construct(
      * @return array{
      *     limit: int,
      *     cursor: string,
+     *     prefix?: string,
      *     includeMetadata: bool,
      *     includeVectors: bool,
      *     includeData: bool,
@@ -25,12 +27,18 @@ public function __construct(
      */
     public function toArray(): array
     {
-        return [
+        $data = [
             'limit' => $this->limit,
             'cursor' => $this->cursor,
             'includeMetadata' => $this->includeMetadata,
             'includeVectors' => $this->includeVectors,
             'includeData' => $this->includeData,
         ];
+
+        if ($this->prefix !== null) {
+            $data['prefix'] = $this->prefix;
+        }
+
+        return $data;
     }
 }
diff --git a/tests/Dense/Operations/RangeVectorsTest.php b/tests/Dense/Operations/RangeVectorsTest.php
index b51372e..0fd397a 100644
--- a/tests/Dense/Operations/RangeVectorsTest.php
+++ b/tests/Dense/Operations/RangeVectorsTest.php
@@ -33,6 +33,22 @@ private function generateUpserts(int $count): array
         return $upserts;
     }
 
+    /**
+     * @return array<VectorUpsert>
+     */
+    private function generatePrefixedUpserts(int $count): array
+    {
+        $upserts = [];
+        for ($i = 0; $i < $count; $i++) {
+            $upserts[] = new VectorUpsert(
+                id: sprintf('prefix-%s', $i),
+                vector: createRandomVector(2),
+            );
+        }
+
+        return $upserts;
+    }
+
     public function test_can_range_vectors(): void
     {
         // Arrange
@@ -157,4 +173,60 @@ public function test_can_range_vectors_using_iterator_can_break_loop(): void
         // Assert
         $this->assertLessThan(1024 * 1024, $memory);
     }
+
+    public function test_can_range_vectors_over_a_prefix(): void
+    {
+        // Arrange
+        $this->namespace->upsertMany([
+            ...$this->generateUpserts(50), // ids: 0, 1, ..., 49
+            ...$this->generatePrefixedUpserts(50), // ids: prefix-0, prefix-1, ..., prefix-49
+        ]);
+        $this->waitForIndex($this->namespace);
+
+        $memory = $this->measureMemory(function () {
+            // Act
+            $results = $this->namespace->range(new VectorRange(limit: 20, prefix: 'prefix-'));
+
+            // Assert
+            $this->assertCount(20, $results);
+
+            // Act
+            $results = $this->namespace->range(new VectorRange(limit: 50, cursor: $results->nextCursor, prefix: 'prefix-'));
+
+            // Assert
+            $this->assertCount(30, $results); // fails here with 50 items
+        });
+
+        // Assert
+        $this->assertLessThan(1024 * 1024, $memory);
+    }
+
+    public function test_can_range_vectors_over_a_prefix_using_iterator(): void
+    {
+        // Arrange
+        $this->namespace->upsertMany([
+            ...$this->generateUpserts(50), // ids: 0, 1, ..., 49
+            ...$this->generatePrefixedUpserts(50), // ids: prefix-0, prefix-1, ..., prefix-49
+        ]);
+        $this->waitForIndex($this->namespace);
+
+        $memory = $this->measureMemory(function () {
+            // Arrange
+            $count = 0;
+
+            // Act
+            $results = $this->namespace->rangeIterator(new VectorRange(limit: 10, prefix: 'prefix-'));
+
+            // Increment count
+            foreach ($results as $result) {
+                $count++;
+            }
+
+            // Assert
+            $this->assertSame(50, $count);
+        });
+
+        // Assert
+        $this->assertLessThan(1024 * 1024, $memory);
+    }
 }