feat: return the distance a vector query ranked by, and let the index answer it

src/Database/Adapter/Postgres.php

@@ -217,7 +217,7 @@ public function createCollection(string $name, array $attributes = [], array $in
  \"_createdAt\" TIMESTAMP(3) DEFAULT NULL,
  \"_updatedAt\" TIMESTAMP(3) DEFAULT NULL,
  " . \implode(' ', $attributeStrings) . "
- _permissions TEXT DEFAULT NULL
+ _permissions JSONB DEFAULT NULL
  );
  ";
 
@@ -226,20 +226,24 @@ public function createCollection(string $name, array $attributes = [], array $in
  $createdIndex = $this->getShortKey("{$namespace}_{$this->tenant}_{$id}_created");
  $updatedIndex = $this->getShortKey("{$namespace}_{$this->tenant}_{$id}_updated");
  $tenantIdIndex = $this->getShortKey("{$namespace}_{$this->tenant}_{$id}_tenant_id");
+ $permissionsIndex = $this->getShortKey("{$namespace}_{$this->tenant}_{$id}_permissions");
  $collection .= "
  CREATE UNIQUE INDEX \"{$uidIndex}\" ON {$this->getSQLTable($id)} (\"_uid\" COLLATE utf8_ci_ai, \"_tenant\");
  CREATE INDEX \"{$createdIndex}\" ON {$this->getSQLTable($id)} (_tenant, \"_createdAt\");
  CREATE INDEX \"{$updatedIndex}\" ON {$this->getSQLTable($id)} (_tenant, \"_updatedAt\");
  CREATE INDEX \"{$tenantIdIndex}\" ON {$this->getSQLTable($id)} (_tenant, _id);
+ CREATE INDEX \"{$permissionsIndex}\" ON {$this->getSQLTable($id)} USING gin (_permissions);
  ";
  } else {
  $uidIndex = $this->getShortKey("{$namespace}_{$id}_uid");
  $createdIndex = $this->getShortKey("{$namespace}_{$id}_created");
  $updatedIndex = $this->getShortKey("{$namespace}_{$id}_updated");
+ $permissionsIndex = $this->getShortKey("{$namespace}_{$id}_permissions");
  $collection .= "
  CREATE UNIQUE INDEX \"{$uidIndex}\" ON {$this->getSQLTable($id)} (\"_uid\" COLLATE utf8_ci_ai);
  CREATE INDEX \"{$createdIndex}\" ON {$this->getSQLTable($id)} (\"_createdAt\");
  CREATE INDEX \"{$updatedIndex}\" ON {$this->getSQLTable($id)} (\"_updatedAt\");
+ CREATE INDEX \"{$permissionsIndex}\" ON {$this->getSQLTable($id)} USING gin (_permissions);
  ";
  }
 
@@ -1755,15 +1759,15 @@ protected function getSQLCondition(Query $query, array &$binds, ?string $forColl
  }
 
  /**
- * Get vector distance calculation for ORDER BY clause
+ * Get the SQL expression measuring distance between a vector attribute and the query vector
  *
  * @param Query $query
  * @param array<string, mixed> $binds
  * @param string $alias
  * @return string|null
  * @throws DatabaseException
  */
- protected function getVectorDistanceOrder(Query $query, array &$binds, string $alias): ?string
+ protected function getSQLVectorDistance(Query $query, array &$binds, string $alias): ?string
  {
  $query->setAttribute($this->getInternalKeyForAttribute($query->getAttribute()));
 
@@ -1785,6 +1789,61 @@ protected function getVectorDistanceOrder(Query $query, array &$binds, string $a
  };
  }
 
+ /**
+ * @param string $distance
+ * @return string
+ */
+ protected function getSQLReadableDistance(string $distance): string
+ {
+ return "{$distance}::text";
+ }
+
+ /**
+ * Match the permission against the copy carried on the row rather than joining the
+ * permissions table.
+ *
+ * Both hold the same fact, written together, but a semi join has to be resolved before
+ * anything can be ordered, which forces the whole collection to be read whenever the
+ * ordering could otherwise have come from an index. Matching on the row leaves the
+ * planner free to cost the permission against the ordering, so a selective permission
+ * drives from the GIN index and a permissive one is a cheap filter over whichever index
+ * the ordering wanted.
+ *
+ * @param string $collection
+ * @param array<string> $roles
+ * @param string $alias
+ * @param string $type
+ * @return string
+ * @throws DatabaseException
+ */
+ protected function getSQLPermissionsCondition(
+ string $collection,
+ array $roles,
+ string $alias,
+ string $type = Database::PERMISSION_READ
+ ): string {
+ if (!\in_array($type, Database::PERMISSIONS)) {
+ throw new DatabaseException('Unknown permission type: ' . $type);
+ }
+
+ $column = "{$this->quote($alias)}.{$this->quote('_permissions')}";
+
+ // Containment rather than jsonb's ?| key operator: PDO reads a lone ? as a positional
+ // placeholder, and doubling it to escape breaks once a named placeholder is repeated,
+ // which the cursor conditions do. Each role is its own @> so the index can answer them
+ // as a BitmapOr; jsonb_exists_any would express it in one call but is not indexable.
+ $permissions = \array_map(
+ fn ($role) => "{$column} @> {$this->getPDO()->quote(\json_encode(["{$type}(\"{$role}\")"]))}::jsonb",
+ $roles
+ );
+
+ if ($permissions === []) {
+ return 'FALSE';
+ }
+
+ return '(' . \implode(' OR ', $permissions) . ')';
+ }
+
  /**
  * @param string $value
  * @return string

src/Database/Adapter/SQL.php

@@ -20,6 +20,8 @@
 
 abstract class SQL extends Adapter
 {
+ protected const VECTOR_DISTANCE_COLUMN = '_distance';
+
  protected mixed $pdo;
 
  /**
@@ -1810,18 +1812,33 @@ abstract protected function getUpsertStatement(
  ): mixed;
 
  /**
- * Get vector distance calculation for ORDER BY clause
+ * Get the SQL expression measuring distance between a vector attribute and the query vector
  *
  * @param Query $query
  * @param array<string, mixed> $binds
  * @param string $alias
  * @return string|null
  */
- protected function getVectorDistanceOrder(Query $query, array &$binds, string $alias): ?string
+ protected function getSQLVectorDistance(Query $query, array &$binds, string $alias): ?string
  {
  return null;
  }
 
+ /**
+ * Render a vector distance expression in a form safe to read back into PHP
+ *
+ * A distance is undefined for a zero vector and can overflow for a large one, so the
+ * expression can evaluate to NaN or infinity. Those cannot survive the trip into a PHP
+ * float, so the value is carried as text and interpreted during hydration.
+ *
+ * @param string $distance
+ * @return string
+ */
+ protected function getSQLReadableDistance(string $distance): string
+ {
+ return $distance;
+ }
+
  /**
  * @param string $value
  * @return string
@@ -3064,18 +3081,17 @@ public function find(Document $collection, array $queries = [], ?int $limit = 25
 
  $sqlWhere = !empty($where) ? 'WHERE ' . implode(' AND ', $where) : '';
 
- // Add vector distance calculations to ORDER BY
- $vectorOrders = [];
+ $vectorDistances = [];
  foreach ($vectorQueries as $query) {
- $vectorOrder = $this->getVectorDistanceOrder($query, $binds, $alias);
- if ($vectorOrder) {
- $vectorOrders[] = $vectorOrder;
+ $vectorDistance = $this->getSQLVectorDistance($query, $binds, $alias);
+ if ($vectorDistance) {
+ $vectorDistances[] = $vectorDistance;
  }
  }
 
- if (!empty($vectorOrders)) {
+ if (!empty($vectorDistances)) {
  // Vector orders should come first for similarity search
- $orders = \array_merge($vectorOrders, $orders);
+ $orders = \array_merge($vectorDistances, $orders);
  }
 
  $sqlOrder = !empty($orders) ? 'ORDER BY ' . implode(', ', $orders) : '';
@@ -3093,8 +3109,15 @@ public function find(Document $collection, array $queries = [], ?int $limit = 25
 
  $selections = $this->getAttributeSelections($queries);
 
+ $projection = $this->getAttributeProjection($selections, $alias);
+
+ if (!empty($vectorDistances)) {
+ $readable = $this->getSQLReadableDistance($vectorDistances[0]);
+ $projection .= ", {$readable} AS {$this->quote(static::VECTOR_DISTANCE_COLUMN)}";
+ }
+
  $sql = "
- SELECT {$this->getAttributeProjection($selections, $alias)}
+ SELECT {$projection}
  FROM {$this->getSQLTable($name)} AS {$this->quote($alias)}
  {$sqlWhere}
  {$sqlOrder}
@@ -3147,6 +3170,11 @@ public function find(Document $collection, array $queries = [], ?int $limit = 25
  $results[$index]['$permissions'] = \json_decode($document['_permissions'] ?? '[]', true);
  unset($results[$index]['_permissions']);
  }
+ if (\array_key_exists(static::VECTOR_DISTANCE_COLUMN, $document)) {
+ $value = $document[static::VECTOR_DISTANCE_COLUMN];
+ $results[$index][Database::VECTOR_DISTANCE] = \is_numeric($value) ? (float)$value : null;
+ unset($results[$index][static::VECTOR_DISTANCE_COLUMN]);
+ }
 
  $results[$index] = new Document($results[$index]);
  }

src/Database/Database.php

@@ -63,6 +63,9 @@ class Database
  // Vector types
  public const VAR_VECTOR = 'vector';
 
+ // Vector query result key
+ public const VECTOR_DISTANCE = '$distance';
+
  // Relationship Types
  public const VAR_RELATIONSHIP = 'relationship';
 
@@ -8559,7 +8562,19 @@ public function find(string $collection, array $queries = [], string $forPermiss
  }
  }
 
- if ($uniqueOrderBy === false) {
+ $vectorSearch = false;
+ foreach ($filters as $filter) {
+ if (\in_array($filter->getMethod(), Query::VECTOR_TYPES)) {
+ $vectorSearch = true;
+ break;
+ }
+ }
+
+ // A vector search is ordered by distance, and a vector index can only answer that one
+ // sort key. Appending a tie break makes the ordering unsatisfiable from the index and
+ // costs a full scan of the collection. The tie break exists to hold a page boundary
+ // still, so it is only owed to a cursor.
+ if ($uniqueOrderBy === false && (!$vectorSearch || !empty($cursor))) {
  $orderAttributes[] = '$sequence';
  }
 

tests/e2e/Adapter/PostgresTest.php

@@ -7,7 +7,11 @@
 use Utopia\Cache\Cache;
 use Utopia\Database\Adapter\Postgres;
 use Utopia\Database\Database;
+use Utopia\Database\Document;
+use Utopia\Database\Helpers\Permission;
+use Utopia\Database\Helpers\Role;
 use Utopia\Database\PDO;
+use Utopia\Database\Query;
 
 class PostgresTest extends Base
 {
@@ -72,4 +76,120 @@ protected function deleteIndex(string $collection, string $index): bool
  return true;
  }
 
+ /**
+ * Reading must be answerable from the row alone. The permissions table holds the same fact,
+ * but reaching it needs a join, and a join has to be resolved before anything can be ordered,
+ * which costs a full read of the collection whenever the ordering could have come from an
+ * index instead.
+ */
+ public function testReadDoesNotTouchThePermissionsTable(): void
+ {
+ $database = $this->getDatabase();
+
+ // no collection level read, so the permission is enforced per document
+ $database->createCollection('permsPlan', permissions: [
+ Permission::create(Role::any()),
+ ], documentSecurity: true);
+
+ $database->createAttribute('permsPlan', 'title', Database::VAR_STRING, 64, true);
+
+ foreach (['visible' => Role::any(), 'hidden' => Role::user('nobody')] as $title => $role) {
+ $database->createDocument('permsPlan', new Document([
+ '$permissions' => [Permission::read($role)],
+ 'title' => $title,
+ ]));
+ }
+
+ $table = $database->getNamespace() . '_permsPlan_perms';
+
+ $scans = function () use ($table): int {
+ self::$pdo->query('SELECT pg_stat_force_next_flush()');
+ self::$pdo->query('SELECT pg_stat_clear_snapshot()');
+
+ $statement = self::$pdo->prepare('SELECT COALESCE(SUM(seq_scan + COALESCE(idx_scan, 0)), 0) FROM pg_stat_user_tables WHERE relname = :table');
+ $statement->execute([':table' => $table]);
+
+ return (int)$statement->fetchColumn();
+ };
+
+ $before = $scans();
+
+ $results = $database->find('permsPlan');
+
+ $this->assertCount(1, $results, 'Only the readable document may come back');
+ $this->assertSame('visible', $results[0]->getAttribute('title'));
+
+ $this->assertSame(
+ $before,
+ $scans(),
+ 'A read must be satisfied from the row, without reaching the permissions table'
+ );
+
+ $database->deleteCollection('permsPlan');
+ }
+
+ /**
+ * A vector search must order by distance alone, because a vector index can answer exactly
+ * one sort key. Adding a second one does not merely make the index look expensive, it makes
+ * it unusable, and the collection is read in full instead.
+ *
+ * Sequential scans are priced out of the session so that the planner falls back to one only
+ * when the index genuinely cannot answer the ordering. That separates a hard block from a
+ * costing preference, and keeps the assertion independent of how many rows are present.
+ */
+ public function testVectorSearchUsesTheIndex(): void
+ {
+ $database = $this->getDatabase();
+
+ $database->createCollection('vectorPlan', permissions: [
+ Permission::read(Role::any()),
+ Permission::create(Role::any()),
+ ], documentSecurity: false);
+
+ $database->createAttribute('vectorPlan', 'embedding', Database::VAR_VECTOR, 3, true);
+ $database->createIndex('vectorPlan', 'idx_cosine', Database::INDEX_HNSW_COSINE, ['embedding']);
+
+ for ($i = 0; $i < 50; $i++) {
+ $database->createDocument('vectorPlan', new Document([
+ '$permissions' => [Permission::read(Role::any())],
+ 'embedding' => [$i / 50, 1 - ($i / 50), 0.0],
+ ]));
+ }
+
+ $index = $database->getNamespace() . '_' . $database->getTenant() . '_vectorPlan_idx_cosine';
+
+ $scans = function () use ($index): int {
+ self::$pdo->query('SELECT pg_stat_force_next_flush()');
+ self::$pdo->query('SELECT pg_stat_clear_snapshot()');
+
+ $statement = self::$pdo->prepare('SELECT COALESCE(SUM(idx_scan), 0) FROM pg_stat_user_indexes WHERE indexrelname = :index');
+ $statement->execute([':index' => $index]);
+
+ return (int)$statement->fetchColumn();
+ };
+
+ $before = $scans();
+
+ self::$pdo->exec('SET enable_seqscan = off');
+
+ try {
+ $results = $database->find('vectorPlan', [
+ Query::vectorCosine('embedding', [1.0, 0.0, 0.0]),
+ Query::limit(10),
+ ]);
+ } finally {
+ self::$pdo->exec('RESET enable_seqscan');
+ }
+
+ $this->assertCount(10, $results);
+ $this->assertEqualsWithDelta(0.0, $results[0]->getAttribute(Database::VECTOR_DISTANCE), 0.001);
+
+ $this->assertGreaterThan(
+ $before,
+ $scans(),
+ 'A vector search must be answerable from the vector index, not by reading the collection'
+ );
+
+ $database->deleteCollection('vectorPlan');
+ }
 }