Added types for Doctrine

Author: ankane

composer.json

@@ -15,16 +15,20 @@
     "autoload": {
         "psr-4": {
             "Pgvector\\": "src/",
+            "Pgvector\\Doctrine\\": "src/doctrine/",
             "Pgvector\\Laravel\\": "src/laravel/"
         }
     },
     "require": {
         "php": ">= 8.1"
     },
     "require-dev": {
+        "doctrine/dbal": "^4",
+        "doctrine/orm": "^3",
         "phpunit/phpunit": "^10",
         "illuminate/database": ">= 10",
-        "laravel/serializable-closure": "^1.3"
+        "laravel/serializable-closure": "^1.3",
+        "symfony/cache": "^7"
     },
     "extra": {
         "laravel": {

src/doctrine/HalfVectorType.php

@@ -0,0 +1,43 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+use Doctrine\DBAL\Types\Type;
+use Doctrine\DBAL\Platforms\AbstractPlatform;
+use Pgvector\HalfVector;
+
+class HalfVectorType extends Type
+{
+    public function getSQLDeclaration(array $fieldDeclaration, AbstractPlatform $platform): string
+    {
+        $length = $fieldDeclaration['length'];
+        return is_null($length) ? 'halfvec' : sprintf('halfvec(%d)', $length);
+    }
+
+    public function convertToPHPValue(mixed $value, AbstractPlatform $platform): ?HalfVector
+    {
+        if (is_null($value)) {
+            return null;
+        }
+
+        return new HalfVector($value);
+    }
+
+    public function convertToDatabaseValue(mixed $value, AbstractPlatform $platform): ?string
+    {
+        if (is_null($value)) {
+            return null;
+        }
+
+        if (!($value instanceof HalfVector)) {
+            $value = new HalfVector($value);
+        }
+
+        return (string) $value;
+    }
+
+    public function getName(): string
+    {
+        return 'halfvec';
+    }
+}

src/doctrine/SparseVectorType.php

@@ -0,0 +1,43 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+use Doctrine\DBAL\Types\Type;
+use Doctrine\DBAL\Platforms\AbstractPlatform;
+use Pgvector\SparseVector;
+
+class SparseVectorType extends Type
+{
+    public function getSQLDeclaration(array $fieldDeclaration, AbstractPlatform $platform): string
+    {
+        $length = $fieldDeclaration['length'];
+        return is_null($length) ? 'sparsevec' : sprintf('sparsevec(%d)', $length);
+    }
+
+    public function convertToPHPValue(mixed $value, AbstractPlatform $platform): ?SparseVector
+    {
+        if (is_null($value)) {
+            return null;
+        }
+
+        return new SparseVector($value);
+    }
+
+    public function convertToDatabaseValue(mixed $value, AbstractPlatform $platform): ?string
+    {
+        if (is_null($value)) {
+            return null;
+        }
+
+        if (!($value instanceof SparseVector)) {
+            $value = new SparseVector($value);
+        }
+
+        return (string) $value;
+    }
+
+    public function getName(): string
+    {
+        return 'sparsevec';
+    }
+}

src/doctrine/VectorType.php

@@ -0,0 +1,43 @@
+<?php
+
+namespace Pgvector\Doctrine;
+
+use Doctrine\DBAL\Types\Type;
+use Doctrine\DBAL\Platforms\AbstractPlatform;
+use Pgvector\Vector;
+
+class VectorType extends Type
+{
+    public function getSQLDeclaration(array $fieldDeclaration, AbstractPlatform $platform): string
+    {
+        $length = $fieldDeclaration['length'];
+        return is_null($length) ? 'vector' : sprintf('vector(%d)', $length);
+    }
+
+    public function convertToPHPValue(mixed $value, AbstractPlatform $platform): ?Vector
+    {
+        if (is_null($value)) {
+            return null;
+        }
+
+        return new Vector($value);
+    }
+
+    public function convertToDatabaseValue(mixed $value, AbstractPlatform $platform): ?string
+    {
+        if (is_null($value)) {
+            return null;
+        }
+
+        if (!($value instanceof Vector)) {
+            $value = new Vector($value);
+        }
+
+        return (string) $value;
+    }
+
+    public function getName(): string
+    {
+        return 'vector';
+    }
+}

tests/DoctrineTest.php

@@ -0,0 +1,60 @@
+<?php
+
+use PHPUnit\Framework\TestCase;
+
+use Doctrine\DBAL\DriverManager;
+use Doctrine\DBAL\Exception\TableNotFoundException;
+use Doctrine\DBAL\Types\Type;
+use Doctrine\ORM\EntityManager;
+use Doctrine\ORM\ORMSetup;
+use Doctrine\ORM\Tools\SchemaTool;
+use Pgvector\HalfVector;
+use Pgvector\SparseVector;
+use Pgvector\Vector;
+
+require_once __DIR__ . '/models/DoctrineItem.php';
+
+final class DoctrineTest extends TestCase
+{
+    public function testTypes()
+    {
+        $config = ORMSetup::createAttributeMetadataConfiguration(
+            paths: [__DIR__ . '/models'],
+            isDevMode: true
+        );
+
+        $connection = DriverManager::getConnection([
+            'driver' => 'pgsql',
+            'dbname' => 'pgvector_php_test'
+        ], $config);
+
+        $entityManager = new EntityManager($connection, $config);
+
+        Type::addType('vector', 'Pgvector\Doctrine\VectorType');
+        Type::addType('halfvec', 'Pgvector\Doctrine\HalfVectorType');
+        Type::addType('sparsevec', 'Pgvector\Doctrine\SparseVectorType');
+
+        $schemaManager = $entityManager->getConnection()->createSchemaManager();
+        try {
+            $schemaManager->dropTable('doctrine_items');
+        } catch (TableNotFoundException $e) {
+            // do nothing
+        }
+
+        $schemaTool = new SchemaTool($entityManager);
+        $schemaTool->createSchema([$entityManager->getClassMetadata('DoctrineItem')]);
+
+        $item = new DoctrineItem();
+        $item->setEmbedding(new Vector([1, 2, 3]));
+        $item->setHalfEmbedding(new HalfVector([4, 5, 6]));
+        $item->setSparseEmbedding(new SparseVector([7, 8, 9]));
+        $entityManager->persist($item);
+        $entityManager->flush();
+
+        $itemRepository = $entityManager->getRepository('DoctrineItem');
+        $item = $itemRepository->find(1);
+        $this->assertEquals([1, 2, 3], $item->getEmbedding()->toArray());
+        $this->assertEquals([4, 5, 6], $item->getHalfEmbedding()->toArray());
+        $this->assertEquals([7, 8, 9], $item->getSparseEmbedding()->toArray());
+    }
+}

tests/models/DoctrineItem.php

@@ -0,0 +1,55 @@
+<?php
+
+use Doctrine\ORM\Mapping as ORM;
+use Pgvector\HalfVector;
+use Pgvector\SparseVector;
+use Pgvector\Vector;
+
+#[ORM\Entity]
+#[ORM\Table(name: 'doctrine_items')]
+class DoctrineItem
+{
+    #[ORM\Id]
+    #[ORM\Column(type: 'integer')]
+    #[ORM\GeneratedValue]
+    private int|null $id = null;
+
+    #[ORM\Column(type: 'vector', length: 3)]
+    private Vector $embedding;
+
+    #[ORM\Column(type: 'halfvec', length: 3)]
+    private HalfVector $halfEmbedding;
+
+    #[ORM\Column(type: 'sparsevec', length: 3)]
+    private SparseVector $sparseEmbedding;
+
+    public function getEmbedding(): Vector
+    {
+        return $this->embedding;
+    }
+
+    public function setEmbedding(Vector $embedding): void
+    {
+        $this->embedding = $embedding;
+    }
+
+    public function getHalfEmbedding(): HalfVector
+    {
+        return $this->halfEmbedding;
+    }
+
+    public function setHalfEmbedding(HalfVector $embedding): void
+    {
+        $this->halfEmbedding = $embedding;
+    }
+
+    public function getSparseEmbedding(): SparseVector
+    {
+        return $this->sparseEmbedding;
+    }
+
+    public function setSparseEmbedding(SparseVector $embedding): void
+    {
+        $this->sparseEmbedding = $embedding;
+    }
+}