Indexing

Author: ankitml

.gitignore

@@ -0,0 +1,6 @@
+.venv/
+
+# Python cache
+__pycache__/
+*.py[cod]
+*$py.class

README.md

@@ -0,0 +1,3 @@
+# sqlalchemy-paradedb
+
+SQLAlchemy query helpers for ParadeDB.

paradedb/__init__.py

@@ -0,0 +1,3 @@
+from . import sqlalchemy
+
+__all__ = ["sqlalchemy"]

paradedb/sqlalchemy/__init__.py

@@ -0,0 +1,3 @@
+from . import facets, indexing, pdb, search, select_with
+
+__all__ = ["facets", "indexing", "pdb", "search", "select_with"]

paradedb/sqlalchemy/_functions.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from sqlalchemy import literal
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.elements import ClauseElement, ColumnElement
+
+
+class PDBFunctionWithNamedArgs(ColumnElement[Any]):
+    """Render pdb.<name>(..., key => value) SQL."""
+
+    inherit_cache = True
+
+    def __init__(
+        self,
+        name: str,
+        positional_args: Sequence[ClauseElement],
+        named_args: Sequence[tuple[str, Any]],
+    ) -> None:
+        self.name = name
+        self.positional_args = tuple(positional_args)
+        self.named_args = tuple(named_args)
+
+
+def _render_named_arg_value(value: Any, compiler, **kw: Any) -> str:
+    if isinstance(value, ClauseElement):
+        return compiler.process(value, **kw)
+    return compiler.process(literal(value), **kw)
+
+
+@compiles(PDBFunctionWithNamedArgs, "postgresql")
+def _compile_pdb_function_with_named_args(element: PDBFunctionWithNamedArgs, compiler, **kw: Any) -> str:
+    positional_sql = [compiler.process(arg, **kw) for arg in element.positional_args]
+    named_sql = [f'{name} => {_render_named_arg_value(value, compiler, **kw)}' for name, value in element.named_args]
+    args_sql = ", ".join([*positional_sql, *named_sql])
+    return f"pdb.{element.name}({args_sql})"
+
+
+@compiles(PDBFunctionWithNamedArgs)
+def _compile_pdb_function_with_named_args_default(element: PDBFunctionWithNamedArgs, compiler, **kw: Any) -> str:
+    raise NotImplementedError("ParadeDB function wrappers are only supported for PostgreSQL dialects")

paradedb/sqlalchemy/_pdb_cast.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from sqlalchemy import literal
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.elements import ClauseElement, ColumnElement
+
+
+class PDBCast(ColumnElement[Any]):
+    """Render an expression cast to a ParadeDB pseudo type, e.g. expr::pdb.boost(2)."""
+
+    inherit_cache = True
+
+    def __init__(self, expr: ClauseElement, type_name: str, args: Sequence[Any] = ()) -> None:
+        self.expr = expr
+        self.type_name = type_name
+        self.args = tuple(args)
+
+
+def _render_cast_arg(arg: Any, compiler, **kw: Any) -> str:
+    if isinstance(arg, ClauseElement):
+        return compiler.process(arg, **kw)
+    if isinstance(arg, bool):
+        return "t" if arg else "f"
+    if isinstance(arg, (int, float)):
+        return str(arg)
+    if isinstance(arg, str):
+        return "'" + arg.replace("'", "''") + "'"
+    return compiler.process(literal(arg), **kw)
+
+
+@compiles(PDBCast, "postgresql")
+def _compile_pdb_cast(element: PDBCast, compiler, **kw: Any) -> str:
+    expr_sql = compiler.process(element.expr, **kw)
+    if element.args:
+        args_sql = ", ".join(_render_cast_arg(arg, compiler, **kw) for arg in element.args)
+        return f"{expr_sql}::pdb.{element.type_name}({args_sql})"
+    return f"{expr_sql}::pdb.{element.type_name}()"
+
+
+@compiles(PDBCast)
+def _compile_pdb_cast_default(element: PDBCast, compiler, **kw: Any) -> str:
+    raise NotImplementedError("ParadeDB casts are only supported for PostgreSQL dialects")

paradedb/sqlalchemy/facets.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+
+def value_count(*, field: str) -> dict[str, dict[str, str]]:
+    return {"value_count": {"field": field}}
+
+
+def avg(*, field: str) -> dict[str, dict[str, str]]:
+    return {"avg": {"field": field}}
+
+
+def terms(*, field: str, size: int | None = None) -> dict[str, dict[str, object]]:
+    payload: dict[str, object] = {"field": field}
+    if size is not None:
+        payload["size"] = size
+    return {"terms": payload}
+
+
+def multi(*aggs: dict[str, object]) -> dict[str, object]:
+    merged: dict[str, object] = {}
+    for agg in aggs:
+        merged.update(agg)
+    return merged

paradedb/sqlalchemy/indexing.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from sqlalchemy import Index, event, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import CompileError
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.elements import ClauseElement, ColumnElement
+from sqlalchemy.sql.visitors import InternalTraversal
+
+
+@dataclass(frozen=True)
+class TokenizerSpec:
+    name: str | None = None
+    options: tuple[tuple[str, Any], ...] = ()
+    raw_sql: str | None = None
+    alias: str | None = None
+
+    def render(self) -> str:
+        if self.raw_sql is not None:
+            return self.raw_sql
+
+        if self.name is None:
+            raise ValueError("tokenizer name is required unless raw_sql is provided")
+
+        if not self.options:
+            return f"pdb.{self.name}()"
+
+        rendered_options = ",".join(f"{key}={_format_option_value(value)}" for key, value in self.options)
+        escaped = rendered_options.replace("'", "''")
+        return f"pdb.{self.name}('{escaped}')"
+
+
+def _format_option_value(value: Any) -> str:
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    return str(value)
+
+
+def _build_spec(name: str, *, alias: str | None = None, **kwargs: Any) -> TokenizerSpec:
+    options: dict[str, Any] = {key: value for key, value in kwargs.items() if value is not None}
+    if alias is not None:
+        options["alias"] = alias
+    return TokenizerSpec(name=name, options=tuple(sorted(options.items())), alias=alias)
+
+
+def unicode(*, alias: str | None = None, lowercase: bool | None = None, stemmer: str | None = None) -> TokenizerSpec:
+    # ParadeDB currently exposes this tokenizer as `unicode_words`.
+    return _build_spec("unicode_words", alias=alias, lowercase=lowercase, stemmer=stemmer)
+
+
+def literal(*, alias: str | None = None) -> TokenizerSpec:
+    return _build_spec("literal", alias=alias)
+
+
+def literal_normalized(*, alias: str | None = None) -> TokenizerSpec:
+    return _build_spec("literal_normalized", alias=alias)
+
+
+def ngram(
+    *,
+    alias: str | None = None,
+    min_gram: int | None = None,
+    max_gram: int | None = None,
+    prefix_only: bool | None = None,
+) -> TokenizerSpec:
+    return _build_spec(
+        "ngram",
+        alias=alias,
+        min_gram=min_gram,
+        max_gram=max_gram,
+        prefix_only=prefix_only,
+    )
+
+
+def raw(sql: str, *, alias: str | None = None) -> TokenizerSpec:
+    return TokenizerSpec(raw_sql=sql, alias=alias)
+
+
+class _TokenizeNamespace:
+    unicode = staticmethod(unicode)
+    literal = staticmethod(literal)
+    literal_normalized = staticmethod(literal_normalized)
+    ngram = staticmethod(ngram)
+    raw = staticmethod(raw)
+
+
+tokenize = _TokenizeNamespace()
+
+
+class BM25Field(ColumnElement[Any]):
+    """Represents a ParadeDB BM25 index field expression."""
+
+    inherit_cache = True
+    _traverse_internals = [
+        ("expr", InternalTraversal.dp_clauseelement),
+        ("tokenizer", InternalTraversal.dp_plain_obj),
+    ]
+
+    def __init__(self, expr: ClauseElement, *, tokenizer: TokenizerSpec | None = None) -> None:
+        self.expr = expr
+        self.tokenizer = tokenizer
+
+    @property
+    def table(self):  # pragma: no cover - SQLAlchemy internals may use this dynamically
+        return getattr(self.expr, "table", None)
+
+
+@compiles(BM25Field, "postgresql")
+def _compile_bm25_field(element: BM25Field, compiler, **kw: Any) -> str:
+    expr_sql = compiler.process(element.expr, **kw)
+    if element.tokenizer is None:
+        return expr_sql
+    return f"({expr_sql}::{element.tokenizer.render()})"
+
+
+@compiles(BM25Field)
+def _compile_bm25_field_default(element: BM25Field, compiler, **kw: Any) -> str:
+    raise CompileError("BM25Field is only supported for PostgreSQL dialects")
+
+
+def _is_bm25_index(index: Index) -> bool:
+    using = index.dialect_options["postgresql"].get("using")
+    return bool(using and str(using).lower() == "bm25")
+
+
+def _bm25_field_name(field: BM25Field) -> str | None:
+    return getattr(getattr(field, "expr", None), "name", None)
+
+
+def validate_bm25_index(index: Index) -> None:
+    if not _is_bm25_index(index):
+        return
+
+    if not index.expressions:
+        raise ValueError("BM25 indexes must include at least one BM25Field")
+
+    if not all(isinstance(expr, BM25Field) for expr in index.expressions):
+        raise ValueError("BM25 indexes must use BM25Field for every indexed field")
+
+    aliases: set[str] = set()
+    for expr in index.expressions:
+        tokenizer = expr.tokenizer
+        if tokenizer is None or tokenizer.alias is None:
+            continue
+        if tokenizer.alias in aliases:
+            raise ValueError(f"Duplicate tokenizer alias '{tokenizer.alias}' in BM25 index")
+        aliases.add(tokenizer.alias)
+
+    with_options = index.dialect_options["postgresql"].get("with") or {}
+    key_field = with_options.get("key_field")
+    if not key_field:
+        raise ValueError("BM25 indexes require postgresql_with={'key_field': '<column>'}")
+
+    field_names = {_bm25_field_name(expr) for expr in index.expressions}
+    if key_field not in field_names:
+        raise ValueError(f"BM25 key_field '{key_field}' must match one of the indexed BM25Field columns")
+
+
+@event.listens_for(Index, "before_create")
+def _validate_bm25_before_create(index: Index, connection, **kw: Any) -> None:
+    validate_bm25_index(index)
+
+
+@dataclass(frozen=True)
+class IndexMeta:
+    index_name: str
+    key_field: str | None
+    fields: tuple[str, ...]
+    aliases: dict[str, str]
+
+
+def describe(engine: Engine, table) -> list[IndexMeta]:
+    query = text(
+        """
+        SELECT indexname, indexdef
+        FROM pg_indexes
+        WHERE schemaname = current_schema()
+          AND tablename = :table_name
+          AND indexdef ILIKE '%USING bm25%'
+        ORDER BY indexname
+        """
+    )
+
+    rows = engine.connect().execute(query, {"table_name": table.name}).fetchall()
+    output: list[IndexMeta] = []
+    for row in rows:
+        indexdef: str = row.indexdef
+        key_field: str | None = None
+        marker = "key_field='"
+        marker_idx = indexdef.find(marker)
+        if marker_idx != -1:
+            key_start = marker_idx + len(marker)
+            key_end = indexdef.find("'", key_start)
+            if key_end != -1:
+                key_field = indexdef[key_start:key_end]
+
+        output.append(
+            IndexMeta(
+                index_name=row.indexname,
+                key_field=key_field,
+                fields=(),
+                aliases={},
+            )
+        )
+    return output