Add pgvector, tsvector/tsquery, and pg_trgm support

piccolo/columns/__init__.py

@@ -27,7 +27,10 @@
     Time,
     Timestamp,
     Timestamptz,
+    Tsquery,
+    Tsvector,
     Varchar,
+    Vector,
 )
 from .combination import And, Or, Where
 from .m2m import M2M
@@ -66,7 +69,10 @@
     "Time",
     "Timestamp",
     "Timestamptz",
+    "Tsquery",
+    "Tsvector",
     "Varchar",
+    "Vector",
     "And",
     "Or",
     "Where",

piccolo/columns/column_types.py

@@ -76,6 +76,8 @@ class Band(Table):
     ArrayAll,
     ArrayAny,
     ArrayNotAny,
+    TrgmSimilar,
+    TSMatch,
 )
 from piccolo.columns.operators.string import Concat
 from piccolo.columns.reference import LazyTableReference
@@ -304,7 +306,26 @@ def get_querystring(
 ###############################################################################
 
 
-class Varchar(Column):
+class TrigramMixin:
+    """
+    Adds pg_trgm operator methods to text column types.
+    Requires the ``pg_trgm`` extension on PostgreSQL.
+    """
+
+    def trigram_similar(self, value: str) -> Where:
+        return Where(
+            column=self,  # type: ignore[arg-type]
+            value=value,
+            operator=TrgmSimilar,
+        )
+
+    def trigram_distance(self, value: str) -> QueryString:
+        return QueryString(  # type: ignore[arg-type]
+            "({} <-> {})", self, value
+        )
+
+
+class Varchar(TrigramMixin, Column):
     """
     Used for storing text when you want to enforce character length limits.
     Uses the ``str`` type for values.
@@ -458,7 +479,7 @@ def __set__(self, obj, value: Union[str, None]):
         obj.__dict__[self._meta.name] = value
 
 
-class Text(Column):
+class Text(TrigramMixin, Column):
     """
     Use when you want to store large strings, and don't want to limit the
     string size. Uses the ``str`` type for values.
@@ -3054,3 +3075,104 @@ def __get__(self, obj, objtype=None):
 
     def __set__(self, obj, value: list[Any]):
         obj.__dict__[self._meta.name] = value
+
+
+###############################################################################
+
+
+class Vector(Column):
+    """
+    Stores fixed-length floating-point vectors for similarity search.
+    Requires the ``pgvector`` extension on PostgreSQL.
+
+    **Example**
+
+    .. code-block:: python
+
+        class Item(Table):
+            embedding = Vector(dimensions=1536)
+
+        >>> await Item(embedding=[0.1, 0.2, 0.3]).save()
+
+    .. note:: Postgres only
+
+    """
+
+    required_extension = "vector"
+    value_type = list
+
+    def __init__(self, dimensions: int = 1536, **kwargs):
+        self.dimensions = dimensions
+        super().__init__(**kwargs)
+        self._meta.params["dimensions"] = dimensions
+
+    @property
+    def column_type(self) -> str:
+        if self._meta.engine_type != "postgres":
+            raise NotImplementedError(
+                "Vector column type is only supported on PostgreSQL."
+            )
+        return f"vector({self.dimensions})"
+
+    def cosine_distance(self, value) -> QueryString:
+        return QueryString("({} <=> {}::vector)", self, value)
+
+    def l2_distance(self, value) -> QueryString:
+        return QueryString("({} <-> {}::vector)", self, value)
+
+    def max_inner_product(self, value) -> QueryString:
+        return QueryString("({} <#> {}::vector)", self, value)
+
+
+class Tsvector(Column):
+    """
+    Stores pre-processed full-text search documents.
+    Requires PostgreSQL.
+
+    **Example**
+
+    .. code-block:: python
+
+        class Article(Table):
+            search_vector = Tsvector()
+
+        >>> await Article.select().where(
+        ...     Article.search_vector.matches(ToTsquery("python"))
+        ... )
+
+    .. note:: Postgres only
+
+    """
+
+    value_type = str
+
+    @property
+    def column_type(self) -> str:
+        if self._meta.engine_type != "postgres":
+            raise NotImplementedError(
+                "Tsvector column type is only supported on PostgreSQL."
+            )
+        return "tsvector"
+
+    def matches(self, query) -> Where:
+        return Where(column=self, value=query, operator=TSMatch)
+
+
+class Tsquery(Column):
+    """
+    Stores a full-text search query.
+    Requires PostgreSQL.
+
+    .. note:: Postgres only
+
+    """
+
+    value_type = str
+
+    @property
+    def column_type(self) -> str:
+        if self._meta.engine_type != "postgres":
+            raise NotImplementedError(
+                "Tsquery column type is only supported on PostgreSQL."
+            )
+        return "tsquery"

piccolo/columns/indexes.py

@@ -11,6 +11,8 @@ class IndexMethod(str, Enum):
     hash = "hash"
     gist = "gist"
     gin = "gin"
+    hnsw = "hnsw"
+    ivfflat = "ivfflat"
 
     def __str__(self):
         return f"{self.__class__.__name__}.{self.name}"

piccolo/columns/operators/__init__.py

@@ -14,6 +14,8 @@
     NotEqual,
     NotIn,
     NotLike,
+    TrgmSimilar,
+    TSMatch,
 )
 from .math import Add, Divide, Multiply, Subtract
 from .string import Concat

piccolo/columns/operators/comparison.py

@@ -68,3 +68,11 @@ class ArrayNotAny(ComparisonOperator):
 
 class ArrayAll(ComparisonOperator):
     template = "{value} = ALL ({name})"
+
+
+class TrgmSimilar(ComparisonOperator):
+    template = "{name} % {value}"
+
+
+class TSMatch(ComparisonOperator):
+    template = "{name} @@ {value}"

piccolo/engine/exceptions.py

@@ -1,2 +1,6 @@
 class TransactionError(Exception):
     pass
+
+
+class ExtensionNotInstalled(Exception):
+    pass

piccolo/engine/extensions.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from piccolo.engine.exceptions import ExtensionNotInstalled
+from piccolo.querystring import QueryString
+
+if TYPE_CHECKING:
+    from piccolo.columns import Column
+    from piccolo.columns.indexes import IndexMethod
+    from piccolo.engine.base import Engine
+
+
+async def check_extensions_for_create(
+    engine: Engine, columns: list[Column]
+) -> None:
+    if engine.engine_type != "postgres":
+        return
+    needed = {
+        ext
+        for col in columns
+        if (ext := getattr(col, "required_extension", None))
+    }
+    await _assert_extensions_installed(engine, needed)
+
+
+async def check_extensions_for_index(
+    engine: Engine,
+    method: IndexMethod,
+    operator_class: str | None,
+) -> None:
+    if engine.engine_type != "postgres":
+        return
+    from piccolo.columns.indexes import IndexMethod as IM
+
+    needed: set[str] = set()
+    if method in (IM.hnsw, IM.ivfflat):
+        needed.add("vector")
+    if operator_class and "trgm" in operator_class:
+        needed.add("pg_trgm")
+    await _assert_extensions_installed(engine, needed)
+
+
+def _encode_vector(v: object) -> str:
+    parts = ",".join(str(x) for x in v)  # type: ignore[attr-defined]
+    return f"[{parts}]"
+
+
+def _decode_vector(s: str) -> list:
+    return [float(x) for x in s.strip("[]").split(",")]
+
+
+async def register_codecs(conn: object) -> None:
+    """Register asyncpg codecs for PostgreSQL extension types."""
+    try:
+        await conn.set_type_codec(  # type: ignore[attr-defined]
+            "vector",
+            encoder=_encode_vector,
+            decoder=_decode_vector,
+            schema="public",
+            format="text",
+        )
+    except ValueError as e:
+        if "unknown type" not in str(e):
+            raise
+
+
+async def _assert_extensions_installed(
+    engine: Engine, extension_names: set[str]
+) -> None:
+    for ext in extension_names:
+        result = await engine.run_querystring(
+            QueryString(
+                "SELECT 1 FROM pg_extension WHERE extname = {}",
+                ext,
+            )
+        )
+        if not result:
+            raise ExtensionNotInstalled(
+                f"The '{ext}' PostgreSQL extension is required but not "
+                f'installed. Run: CREATE EXTENSION "{ext}";'
+            )

piccolo/engine/postgres.py

@@ -15,6 +15,7 @@
     validate_savepoint_name,
 )
 from piccolo.engine.exceptions import TransactionError
+from piccolo.engine.extensions import register_codecs
 from piccolo.query.base import DDL, Query
 from piccolo.querystring import QueryString
 from piccolo.utils.lazy_loader import LazyLoader
@@ -465,7 +466,9 @@ async def start_connection_pool(self, **kwargs) -> None:
         else:
             config = dict(self.config)
             config.update(**kwargs)
-            self.pool = await asyncpg.create_pool(**config)
+            self.pool = await asyncpg.create_pool(
+                **config, init=register_codecs
+            )
 
     async def close_connection_pool(self) -> None:
         if self.pool:
@@ -480,7 +483,9 @@ async def get_new_connection(self) -> Connection:
         """
         Returns a new connection - doesn't retrieve it from the pool.
         """
-        return await asyncpg.connect(**self.config)
+        conn = await asyncpg.connect(**self.config)
+        await register_codecs(conn)
+        return conn
 
     ###########################################################################
 

piccolo/query/functions/__init__.py

@@ -27,7 +27,20 @@
     Replace,
     Reverse,
     Rtrim,
+    Similarity,
     Upper,
+    WordSimilarity,
+)
+from .text_search import (
+    PhrasetoTsquery,
+    PlaintoTsquery,
+    Setweight,
+    ToTsquery,
+    ToTsvector,
+    TsHeadline,
+    TsRank,
+    TsRankCd,
+    WebsearchToTsquery,
 )
 from .type_conversion import Cast
 
@@ -48,7 +61,6 @@
     "Count",
     "Day",
     "Extract",
-    "Extract",
     "Floor",
     "Hour",
     "Length",
@@ -58,13 +70,24 @@
     "Min",
     "Month",
     "NullIf",
+    "PhrasetoTsquery",
+    "PlaintoTsquery",
     "Replace",
     "Reverse",
     "Round",
     "Rtrim",
     "Second",
+    "Setweight",
+    "Similarity",
     "Strftime",
     "Sum",
+    "ToTsquery",
+    "ToTsvector",
+    "TsHeadline",
+    "TsRank",
+    "TsRankCd",
     "Upper",
+    "WebsearchToTsquery",
+    "WordSimilarity",
     "Year",
 )