Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/Iterators/VectorRangeIterator.php
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
10 changes: 9 additions & 1 deletion src/VectorRange.php
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -18,19 +19,26 @@ public function __construct(
* @return array{
* limit: int,
* cursor: string,
* prefix?: string,
* includeMetadata: bool,
* includeVectors: bool,
* includeData: bool,
* }
*/
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;
}
}
72 changes: 72 additions & 0 deletions tests/Dense/Operations/RangeVectorsTest.php
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
}