feat: Various improvements and clean ups (#30)

Author: isaacvando

.github/workflows/ci.yml

@@ -50,6 +50,9 @@ jobs:
       - name: Type check
         run: uv run --extra dev mypy paradedb
 
+      - name: Check API coverage
+        run: uv run scripts/check_api_coverage.py
+
   matrix-tests:
     name: Python ${{ matrix.python-version }}
     runs-on: ubuntu-latest

AGENTS.md

@@ -0,0 +1 @@
+CLAUDE.md

CLAUDE.md

@@ -2,24 +2,17 @@
 
 ## Unit Tests
 
-Run unit tests (no database required) before every push — this is what CI checks first:
+Run the tests to verify every change:
 
 ```bash
 .venv/bin/python -m pytest -m "not integration"
-```
-
-## Integration Tests
-
-Integration tests require a running ParadeDB/Postgres instance. Use the provided script to start the container and set up the correct DSN:
-
-```bash
-PARADEDB_PASSWORD=postgres scripts/run_integration_tests.sh
+scripts/run_integration_tests.sh
 ```
 
 To run a subset of tests, pass pytest selectors:
 
 ```bash
-PARADEDB_PASSWORD=postgres scripts/run_integration_tests.sh tests/integration/test_indexing_integration.py::test_bm25_partial_index_generates_where_clause
+scripts/run_integration_tests.sh tests/integration/test_indexing_integration.py::test_bm25_partial_index_generates_where_clause
 ```
 
 The script sets `PARADEDB_TEST_DSN` and `DATABASE_URL` automatically. The default container name is `paradedb-sqlalchemy-integration` on port `5443`.

README.md

@@ -92,7 +92,7 @@ base = (
     .limit(10)
 )
 
-stmt, facet_plan = facets.with_rows(
+stmt = facets.with_rows(
     base,
     agg=facets.multi(
         facets.value_count(field="id"),
@@ -103,7 +103,7 @@ stmt, facet_plan = facets.with_rows(
 
 with Session(engine) as session:
     rows = session.execute(stmt).all()
-    facet_payload = facet_plan.extract(rows)
+    facet_payload = facets.extract(rows)
 ```
 
 ## Search Patterns
@@ -135,11 +135,7 @@ search.more_like_this(Product.id, document_id=1, fields=["description"])
 from sqlalchemy import select
 from paradedb.sqlalchemy import search
 
-prox = search.prox_array("running").near(
-    search.prox_regex("sho.*"),
-    distance=1,
-    ordered=True,
-)
+prox = search.prox_array("running").within(1, search.prox_regex("sho.*"), ordered=True)
 stmt = select(Product.id).where(search.proximity(Product.description, prox))
 ```
 

api.json

@@ -1,54 +1,57 @@
 {
   "operators": {
-    "OP_SEARCH":        "@@@",
-    "OP_AND":           "&&&",
-    "OP_OR":            "|||",
-    "OP_PHRASE":        "###",
-    "OP_PROXIMITY":     "##",
+    "OP_SEARCH":       "@@@",
+    "OP_AND":          "&&&",
+    "OP_OR":           "|||",
+    "OP_PHRASE":       "###",
+    "OP_PROXIMITY":    "##",
     "OP_PROXIMITY_ORD": "##>",
-    "OP_TERM":          "==="
+    "OP_TERM":         "==="
   },
 
   "functions": {
-    "FN_ALL":              "pdb.all",
-    "FN_PARSE":            "pdb.parse",
-    "FN_PHRASE_PREFIX":    "pdb.phrase_prefix",
-    "FN_REGEX_PHRASE":     "pdb.regex_phrase",
-    "FN_RANGE_TERM":       "pdb.range_term",
-    "FN_REGEX":            "pdb.regex",
-    "FN_PROX_REGEX":       "pdb.prox_regex",
-    "FN_PROX_ARRAY":       "pdb.prox_array",
-    "FN_MORE_LIKE_THIS":   "pdb.more_like_this",
-    "FN_EXISTS":           "pdb.exists",
+    "FN_ALL":                "pdb.all",
+    "FN_PARSE":              "pdb.parse",
+    "FN_PHRASE_PREFIX":      "pdb.phrase_prefix",
+    "FN_REGEX_PHRASE":       "pdb.regex_phrase",
+    "FN_RANGE_TERM":         "pdb.range_term",
+    "FN_REGEX":              "pdb.regex",
+    "FN_PROX_REGEX":         "pdb.prox_regex",
+    "FN_PROX_ARRAY":         "pdb.prox_array",
+    "FN_MORE_LIKE_THIS":     "pdb.more_like_this",
+    "FN_EXISTS":             "pdb.exists",
 
-    "FN_SCORE":            "pdb.score",
-    "FN_SNIPPET":          "pdb.snippet",
-    "FN_SNIPPETS":         "pdb.snippets",
-    "FN_SNIPPET_POSITIONS": "pdb.snippet_positions",
-    "FN_AGG":              "pdb.agg",
+    "FN_SCORE":              "pdb.score",
+    "FN_SNIPPET":            "pdb.snippet",
+    "FN_SNIPPETS":           "pdb.snippets",
+    "FN_SNIPPET_POSITIONS":  "pdb.snippet_positions",
+    "FN_AGG":                "pdb.agg",
 
-    "FN_INDEXES":          "pdb.indexes",
-    "FN_INDEX_SEGMENTS":   "pdb.index_segments",
-    "FN_VERIFY_INDEX":     "pdb.verify_index",
-    "FN_VERIFY_ALL_INDEXES": "pdb.verify_all_indexes"
+    "FN_INDEX_SEGMENTS":     "pdb.index_segments",
+    "FN_INDEXES":            "pdb.indexes",
+    "FN_VERIFY_ALL_INDEXES": "pdb.verify_all_indexes",
+    "FN_VERIFY_INDEX":       "pdb.verify_index"
   },
 
   "types": {
-    "PDB_TYPE_BOOST":            "pdb.boost",
-    "PDB_TYPE_CONST":            "pdb.const",
-    "PDB_TYPE_FUZZY":            "pdb.fuzzy",
-    "PDB_TYPE_SLOP":             "pdb.slop",
-    "PDB_TYPE_QUERY":            "pdb.query",
-    "PDB_TYPE_PROXIMITYCLAUSE":  "pdb.proximityclause",
+    "PDB_TYPE_BOOST": "pdb.boost",
+    "PDB_TYPE_CONST": "pdb.const",
+    "PDB_TYPE_FUZZY": "pdb.fuzzy",
+    "PDB_TYPE_SLOP":  "pdb.slop",
+    "PDB_TYPE_QUERY": "pdb.query",
 
-    "PDB_TYPE_TOKENIZER_ICU":                "pdb.icu",
+    "PDB_TYPE_TOKENIZER_ALIAS":              "pdb.alias",
+    "PDB_TYPE_TOKENIZER_CHINESE_COMPATIBLE": "pdb.chinese_compatible",
+    "PDB_TYPE_TOKENIZER_JIEBA":              "pdb.jieba",
     "PDB_TYPE_TOKENIZER_LINDERA":            "pdb.lindera",
     "PDB_TYPE_TOKENIZER_LITERAL":            "pdb.literal",
     "PDB_TYPE_TOKENIZER_LITERAL_NORMALIZED": "pdb.literal_normalized",
     "PDB_TYPE_TOKENIZER_NGRAM":              "pdb.ngram",
     "PDB_TYPE_TOKENIZER_REGEX":              "pdb.regex_pattern",
     "PDB_TYPE_TOKENIZER_SIMPLE":             "pdb.simple",
+    "PDB_TYPE_TOKENIZER_SOURCE_CODE":        "pdb.source_code",
     "PDB_TYPE_TOKENIZER_UNICODE_WORDS":      "pdb.unicode_words",
-    "PDB_TYPE_TOKENIZER_WHITESPACE":         "pdb.whitespace"
+    "PDB_TYPE_TOKENIZER_WHITESPACE":         "pdb.whitespace",
+    "PDB_TYPE_TOKENIZER_ICU":                "pdb.icu"
   }
 }

apiignore.json

@@ -207,8 +207,6 @@
 
   "types": [
     "pdb.alias",
-    "pdb.chinese_compatible",
-    "pdb.jieba",
-    "pdb.source_code"
+    "pdb.proximityclause"
   ]
 }

paradedb/__init__.py

@@ -7,7 +7,7 @@
 )
 from .sqlalchemy.facets import with_rows
 from .sqlalchemy.indexing import BM25Field, assert_indexed, describe, tokenize
-from .sqlalchemy.pdb import agg, score, snippet, snippet_positions, snippets
+from .sqlalchemy.pdb import agg, alias, score, snippet, snippet_positions, snippets
 from .sqlalchemy.search import (
     ProximityExpr,
     all,
@@ -30,6 +30,7 @@
     "BM25Field",
     "ProximityExpr",
     "agg",
+    "alias",
     "all",
     "assert_indexed",
     "describe",

paradedb/sqlalchemy/facets.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
 from collections.abc import Sequence
 from typing import Any
 
@@ -102,20 +101,16 @@ def ensure_operator(stmt: Select, *, key_field: ColumnElement) -> Select:
     return stmt.where(search.all(key_field))
 
 
-@dataclass(frozen=True)
-class FacetPlan:
-    label: str = "facets"
-
-    def extract(self, rows: list[object]) -> Any | None:
-        if not rows:
-            return None
-        first = rows[0]
-        mapping = getattr(first, "_mapping", None)
-        if mapping is not None and self.label in mapping:
-            return mapping[self.label]
-        if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
-            return first[-1]
+def extract(rows: Sequence[object], *, label: str = "facets") -> Any | None:
+    if not rows:
         return None
+    first = rows[0]
+    mapping = getattr(first, "_mapping", None)
+    if mapping is not None and label in mapping:
+        return mapping[label]
+    if isinstance(first, Sequence) and not isinstance(first, (str, bytes)):
+        return first[-1]
+    return None
 
 
 def with_rows(
@@ -125,7 +120,7 @@ def with_rows(
     key_field: ColumnElement,
     label: str = "facets",
     ensure_predicate: bool = True,
-) -> tuple[Select, FacetPlan]:
+) -> Select:
     if not has_order_by(base_stmt):
         raise FacetRequiresOrderByError("with_rows requires ORDER BY")
     if not has_limit(base_stmt):
@@ -136,4 +131,4 @@ def with_rows(
         raise FacetRequiresParadeDBPredicateError("with_rows requires a ParadeDB predicate")
 
     stmt = stmt.add_columns(pdb.agg(agg).over().label(label))
-    return stmt, FacetPlan(label=label)
+    return stmt

paradedb/sqlalchemy/indexing.py

@@ -12,6 +12,7 @@
 from sqlalchemy.sql.visitors import InternalTraversal
 
 from ._select_introspection import has_limit, has_order_by
+from ._pdb_cast import PDBCast
 from .errors import (
     DuplicateTokenizerAliasError,
     FieldNotIndexedError,
@@ -198,6 +199,42 @@ def icu(
     )
 
 
+def chinese_compatible(
+    *,
+    alias: str | None = None,
+    args: Sequence[Any] | None = None,
+    named_args: Mapping[str, Any] | None = None,
+    filters: Sequence[str] | None = None,
+    stemmer: str | None = None,
+) -> TokenizerSpec:
+    return _build_spec(
+        "chinese_compatible",
+        alias=alias,
+        args=args,
+        named_args=named_args,
+        filters=filters,
+        stemmer=stemmer,
+    )
+
+
+def jieba(
+    *,
+    alias: str | None = None,
+    args: Sequence[Any] | None = None,
+    named_args: Mapping[str, Any] | None = None,
+    filters: Sequence[str] | None = None,
+    stemmer: str | None = None,
+) -> TokenizerSpec:
+    return _build_spec(
+        "jieba",
+        alias=alias,
+        args=args,
+        named_args=named_args,
+        filters=filters,
+        stemmer=stemmer,
+    )
+
+
 def literal(*, alias: str | None = None) -> TokenizerSpec:
     return _build_spec("literal", alias=alias)
 
@@ -286,6 +323,24 @@ def regex_pattern(
     )
 
 
+def source_code(
+    *,
+    alias: str | None = None,
+    args: Sequence[Any] | None = None,
+    named_args: Mapping[str, Any] | None = None,
+    filters: Sequence[str] | None = None,
+    stemmer: str | None = None,
+) -> TokenizerSpec:
+    return _build_spec(
+        "source_code",
+        alias=alias,
+        args=args,
+        named_args=named_args,
+        filters=filters,
+        stemmer=stemmer,
+    )
+
+
 def raw(sql: str, *, alias: str | None = None) -> TokenizerSpec:
     return TokenizerSpec(raw_sql=sql, alias=alias)
 
@@ -340,11 +395,14 @@ class _TokenizeNamespace:
     simple = staticmethod(simple)
     whitespace = staticmethod(whitespace)
     icu = staticmethod(icu)
+    chinese_compatible = staticmethod(chinese_compatible)
+    jieba = staticmethod(jieba)
     literal = staticmethod(literal)
     literal_normalized = staticmethod(literal_normalized)
     ngram = staticmethod(ngram)
     lindera = staticmethod(lindera)
     regex_pattern = staticmethod(regex_pattern)
+    source_code = staticmethod(source_code)
     raw = staticmethod(raw)
     custom = staticmethod(custom)
     from_config = staticmethod(from_config)
@@ -375,6 +433,8 @@ def table(self):  # pragma: no cover - SQLAlchemy internals may use this dynamic
 def _compile_bm25_field(element: BM25Field, compiler, **kw: Any) -> str:
     expr_sql = compiler.process(element.expr, **kw)
     if element.tokenizer is None:
+        if isinstance(element.expr, PDBCast) or _bm25_field_name(element) is None:
+            return f"({expr_sql})"
         return expr_sql
     return f"(({expr_sql})::{element.tokenizer.render()})"
 

paradedb/sqlalchemy/pdb.py

@@ -3,19 +3,28 @@
 import json
 from typing import Any
 
-from sqlalchemy import cast, func, literal
-from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy import Text, func, literal, literal_column
 from sqlalchemy.sql.elements import ClauseElement, ColumnElement
 
 from .errors import InvalidArgumentError
 from ._functions import PDBFunctionWithNamedArgs
+from ._pdb_cast import PDBCast
 from .validation import require_non_empty_string, require_non_negative, require_positive
 
 
+def _inline_string_literal(value: str) -> ClauseElement:
+    return literal_column("'" + value.replace("'", "''") + "'", Text())
+
+
 def score(field: ColumnElement) -> ClauseElement:
     return func.pdb.score(field)
 
 
+def alias(field: ColumnElement, name: str) -> ClauseElement:
+    require_non_empty_string(name, field_name="name")
+    return PDBCast(field.self_group(), "alias", (name,))
+
+
 def snippet(
     field: ColumnElement,
     *,
@@ -95,7 +104,7 @@ def agg(spec: dict[str, Any], *, approximate: bool | None = None) -> ClauseEleme
     if not isinstance(spec, dict) or not spec:
         raise InvalidArgumentError("spec must be a non-empty dict")
     payload = json.dumps(spec, separators=(",", ":"), sort_keys=True)
-    payload_expr = cast(literal(payload), JSONB)
+    payload_expr = _inline_string_literal(payload)
     if approximate is None:
         return func.pdb.agg(payload_expr)
     # pdb.agg() takes an optional second positional boolean: true = exact (default),