feat: add ability to access fetched vectors by their id

Author: heyjorgedev

src/Operations/FetchVectorsOperation.php

@@ -11,6 +11,7 @@
 use Upstash\Vector\Transporter\TransporterResponse;
 use Upstash\Vector\VectorFetch;
 use Upstash\Vector\VectorFetchResult;
+use Upstash\Vector\VectorMatch;
 
 /**
  * @internal
@@ -58,7 +59,11 @@ private function transformResponse(TransporterResponse $response): VectorFetchRe
         $result = array_filter($result);
 
         $results = array_map(fn (array $result) => $this->mapVectorMatch($result), $result);
+        $keys = array_map(fn (VectorMatch $match) => $match->getIdentifier(), $results);
 
-        return new VectorFetchResult($results);
+        // Key-by vector ID
+        $results = array_combine($keys, $results);
+
+        return new VectorFetchResult($keys, $results);
     }
 }

src/VectorFetchResult.php

@@ -10,25 +10,38 @@
 
 /**
  * @implements IteratorAggregate<int,VectorMatch>
- * @implements ArrayAccess<int,VectorMatch>
+ * @implements ArrayAccess<int,VectorMatch|null>
  */
 final readonly class VectorFetchResult implements ArrayAccess, Countable, IteratorAggregate
 {
     /**
      * @param  array<VectorMatch>  $results
      */
     public function __construct(
+        private array $keys,
         public array $results,
     ) {}
 
     public function offsetExists(mixed $offset): bool
     {
+        if (is_int($offset)) {
+            $offset = $this->keys[$offset] ?? $offset;
+        }
+
         return isset($this->results[$offset]);
     }
 
-    public function offsetGet(mixed $offset): VectorMatch
+    public function offsetGet(mixed $offset): ?VectorMatch
     {
-        return $this->results[$offset];
+        if (is_int($offset)) {
+            $offset = $this->keys[$offset] ?? $offset;
+        }
+
+        if ($this->offsetExists($offset)) {
+            return $this->results[$offset];
+        }
+
+        return null;
     }
 
     public function offsetSet(mixed $offset, mixed $value): void {}

tests/Dense/Operations/FetchVectorsOperationTest.php

@@ -33,7 +33,30 @@ public function test_can_fetch_vectors(): void
 
         // Assert
         $this->assertCount(2, $results);
-        $this->assertCount(2, $results[0]->vector);
+        $this->assertNotNull($results[0]->vector);
+        $this->assertCount(2, $results[0]->vector); // O(n)
+    }
+
+    public function test_can_fetch_vectors_and_reference_them_by_vector_id(): void
+    {
+        // Arrange
+        $this->namespace->upsertMany([
+            new VectorUpsert(id: '1', vector: createRandomVector(2)),
+            new VectorUpsert(id: '2', vector: createRandomVector(2)),
+            new VectorUpsert(id: '3', vector: createRandomVector(2)),
+        ]);
+        $this->waitForIndex($this->namespace);
+
+        // Act
+        $results = $this->namespace->fetch(new VectorFetch(
+            ids: ['1', '2'],
+            includeVectors: true,
+        ));
+
+        // Assert
+        $this->assertCount(2, $results);
+        $this->assertNotNull($results['1']->vector);
+        $this->assertCount(2, $results['1']->vector); // O(n)
     }
 
     public function test_can_fetch_vectors_that_dont_exist(): void
@@ -47,4 +70,17 @@ public function test_can_fetch_vectors_that_dont_exist(): void
         // Assert
         $this->assertCount(0, $results);
     }
+
+    public function test_can_fetch_vectors_that_dont_exist_and_get_null(): void
+    {
+        // Act
+        $results = $this->namespace->fetch(new VectorFetch(
+            ids: ['1', '2'],
+            includeVectors: true,
+        ));
+
+        // Assert
+        $this->assertNull($results[0]);
+        $this->assertNull($results['1']);
+    }
 }