Added distance functions for Doctrine

Author: ankane

src/doctrine/CosineDistance.php

@@ -0,0 +1,11 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+class CosineDistance extends DistanceNode
+{
+    protected function getOp(): string
+    {
+        return '<=>';
+    }
+}

src/doctrine/DistanceNode.php

@@ -0,0 +1,36 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+use Doctrine\ORM\Query\AST\Functions\FunctionNode;
+use Doctrine\ORM\Query\Parser;
+use Doctrine\ORM\Query\SqlWalker;
+use Doctrine\ORM\Query\TokenType;
+
+abstract class DistanceNode extends FunctionNode
+{
+    public $left;
+    public $right;
+
+    abstract protected function getOp(): string;
+
+    public function parse(Parser $parser): void
+    {
+        $parser->match(TokenType::T_IDENTIFIER);
+        $parser->match(TokenType::T_OPEN_PARENTHESIS);
+        $this->left = $parser->ArithmeticPrimary();
+        $parser->match(TokenType::T_COMMA);
+        $this->right = $parser->ArithmeticPrimary();
+        $parser->match(TokenType::T_CLOSE_PARENTHESIS);
+    }
+
+    public function getSql(SqlWalker $sqlWalker): string
+    {
+        return sprintf(
+            '(%s %s %s)',
+            $this->left->dispatch($sqlWalker),
+            $this->getOp(),
+            $this->right->dispatch($sqlWalker),
+        );
+    }
+}

src/doctrine/L1Distance.php

@@ -0,0 +1,11 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+class L1Distance extends DistanceNode
+{
+    protected function getOp(): string
+    {
+        return '<+>';
+    }
+}

src/doctrine/L2Distance.php

@@ -0,0 +1,11 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+class L2Distance extends DistanceNode
+{
+    protected function getOp(): string
+    {
+        return '<->';
+    }
+}

src/doctrine/MaxInnerProduct.php

@@ -0,0 +1,11 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+class MaxInnerProduct extends DistanceNode
+{
+    protected function getOp(): string
+    {
+        return '<#>';
+    }
+}

tests/DoctrineTest.php

@@ -25,6 +25,10 @@ public static function setUpBeforeClass(): void
             paths: [__DIR__ . '/models'],
             isDevMode: true
         );
+        $config->addCustomNumericFunction('l2_distance', 'Pgvector\Doctrine\L2Distance');
+        $config->addCustomNumericFunction('max_inner_product', 'Pgvector\Doctrine\MaxInnerProduct');
+        $config->addCustomNumericFunction('cosine_distance', 'Pgvector\Doctrine\CosineDistance');
+        $config->addCustomNumericFunction('l1_distance', 'Pgvector\Doctrine\L1Distance');
 
         $connection = DriverManager::getConnection([
             'driver' => 'pgsql',
@@ -84,10 +88,9 @@ public function testTypes()
     public function testVectorL2Distance()
     {
         $this->createItems();
-        $rsm = new ResultSetMappingBuilder(self::$em);
-        $rsm->addRootEntityFromClassMetadata('DoctrineItem', 'i');
-        $neighbors = self::$em->createNativeQuery('SELECT * FROM doctrine_items i ORDER BY embedding <-> ? LIMIT 5', $rsm)
+        $neighbors = self::$em->createQuery('SELECT i FROM DoctrineItem i ORDER BY l2_distance(i.embedding, ?1)')
             ->setParameter(1, new Vector([1, 1, 1]))
+            ->setMaxResults(5)
             ->getResult();
         $this->assertEquals([1, 3, 2], array_map(fn ($v) => $v->getId(), $neighbors));
         $this->assertEquals([[1, 1, 1], [1, 1, 2], [2, 2, 2]], array_map(fn ($v) => $v->getEmbedding()->toArray(), $neighbors));
@@ -96,32 +99,29 @@ public function testVectorL2Distance()
     public function testVectorMaxInnerProduct()
     {
         $this->createItems();
-        $rsm = new ResultSetMappingBuilder(self::$em);
-        $rsm->addRootEntityFromClassMetadata('DoctrineItem', 'i');
-        $neighbors = self::$em->createNativeQuery('SELECT * FROM doctrine_items i ORDER BY embedding <#> ? LIMIT 5', $rsm)
+        $neighbors = self::$em->createQuery('SELECT i FROM DoctrineItem i ORDER BY max_inner_product(i.embedding, ?1)')
             ->setParameter(1, new Vector([1, 1, 1]))
+            ->setMaxResults(5)
             ->getResult();
         $this->assertEquals([2, 3, 1], array_map(fn ($v) => $v->getId(), $neighbors));
     }
 
     public function testVectorCosineDistance()
     {
         $this->createItems();
-        $rsm = new ResultSetMappingBuilder(self::$em);
-        $rsm->addRootEntityFromClassMetadata('DoctrineItem', 'i');
-        $neighbors = self::$em->createNativeQuery('SELECT * FROM doctrine_items i ORDER BY embedding <=> ? LIMIT 5', $rsm)
+        $neighbors = self::$em->createQuery('SELECT i FROM DoctrineItem i ORDER BY cosine_distance(i.embedding, ?1)')
             ->setParameter(1, new Vector([1, 1, 1]))
+            ->setMaxResults(5)
             ->getResult();
         $this->assertEquals([1, 2, 3], array_map(fn ($v) => $v->getId(), $neighbors));
     }
 
     public function testVectorL1Distance()
     {
         $this->createItems();
-        $rsm = new ResultSetMappingBuilder(self::$em);
-        $rsm->addRootEntityFromClassMetadata('DoctrineItem', 'i');
-        $neighbors = self::$em->createNativeQuery('SELECT * FROM doctrine_items i ORDER BY embedding <+> ? LIMIT 5', $rsm)
+        $neighbors = self::$em->createQuery('SELECT i FROM DoctrineItem i ORDER BY l1_distance(i.embedding, ?1)')
             ->setParameter(1, new Vector([1, 1, 1]))
+            ->setMaxResults(5)
             ->getResult();
         $this->assertEquals([1, 3, 2], array_map(fn ($v) => $v->getId(), $neighbors));
     }