Added vector distance support to postgres compiler

Author: freost

src/mako/database/query/compilers/Postgres.php

@@ -7,11 +7,14 @@
 
 namespace mako\database\query\compilers;
 
+use mako\database\query\VectorMetric;
 use Override;
 
 use function array_pop;
 use function implode;
+use function is_array;
 use function is_numeric;
+use function json_encode;
 use function str_replace;
 
 /**
@@ -55,6 +58,22 @@ protected function buildJsonSet(string $column, array $segments, string $param):
 		return $column . " = JSONB_SET({$column}, '{" . str_replace("'", "''", implode(',', $segments)) . "}', '{$param}')";
 	}
 
+	/**
+	 * {@inheritDoc}
+	 */
+	#[Override]
+	public function whereVectorDistance(array $where): string
+	{
+		$vector = is_array($where['vector']) ? json_encode($where['vector']) : $where['vector'];
+
+		$operator = match ($where['metric']) {
+			VectorMetric::COSINE => '<=>',
+			VectorMetric::EUCLIDEAN => '<->',
+		};
+
+		return "{$this->column($where['column'], false)} {$operator} {$this->param($vector)} <= {$this->param($where['distance'])}";
+	}
+
 	/**
 	 * {@inheritDoc}
 	 */

tests/unit/database/query/compilers/PostgresCompilerTest.php

@@ -11,6 +11,7 @@
 use mako\database\query\compilers\Postgres as PostgresCompiler;
 use mako\database\query\helpers\HelperInterface;
 use mako\database\query\Query;
+use mako\database\query\VectorMetric;
 use mako\tests\TestCase;
 use Mockery;
 use Mockery\MockInterface;
@@ -204,6 +205,36 @@ public function testUpdateWithJSONColumn(): void
 		$this->assertEquals([1], $query['params']);
 	}
 
+	/**
+	 *
+	 */
+	public function testBasicCosineWhereVectorDistance(): void
+	{
+		$query = $this->getBuilder();
+
+		$query = $query->table('foobar')
+		->whereVectorDistance('embedding', [1, 2, 3, 4, 5], maxDistance: 0.5)
+		->getCompiler()->select();
+
+		$this->assertEquals('SELECT * FROM "foobar" WHERE "embedding" <=> ? <= ?', $query['sql']);
+		$this->assertEquals(['[1,2,3,4,5]', 0.5], $query['params']);
+	}
+
+	/**
+	 *
+	 */
+	public function testBasicEuclidianWhereVectorDistance(): void
+	{
+		$query = $this->getBuilder();
+
+		$query = $query->table('foobar')
+		->whereVectorDistance('embedding', [1, 2, 3, 4, 5], maxDistance: 0.5, vectorMetric: VectorMetric::EUCLIDEAN)
+		->getCompiler()->select();
+
+		$this->assertEquals('SELECT * FROM "foobar" WHERE "embedding" <-> ? <= ?', $query['sql']);
+		$this->assertEquals(['[1,2,3,4,5]', 0.5], $query['params']);
+	}
+
 	/**
 	 *
 	 */