Merge pull request #1 from paradedb/ankitml/querying

feat: bulk of querying

Author: ankitml

paradedb/sqlalchemy/__init__.py

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

paradedb/sqlalchemy/_pdb_cast.py

@@ -37,7 +37,7 @@ def _compile_pdb_cast(element: PDBCast, compiler, **kw: Any) -> str:
     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}()"
+    return f"{expr_sql}::pdb.{element.type_name}"
 
 
 @compiles(PDBCast)

paradedb/sqlalchemy/alembic.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from alembic.operations import Operations
+from alembic.operations.ops import MigrateOperation
+
+
+def _quote_ident(name: str) -> str:
+    return '"' + name.replace('"', '""') + '"'
+
+
+@Operations.register_operation("create_bm25_index")
+class CreateBM25IndexOp(MigrateOperation):
+    def __init__(self, index_name: str, table_name: str, fields: list[str], key_field: str) -> None:
+        self.index_name = index_name
+        self.table_name = table_name
+        self.fields = fields
+        self.key_field = key_field
+
+    @classmethod
+    def create_bm25_index(
+        cls,
+        operations: Operations,
+        index_name: str,
+        table_name: str,
+        fields: list[str],
+        *,
+        key_field: str,
+    ) -> MigrateOperation:
+        return operations.invoke(cls(index_name, table_name, fields, key_field))
+
+
+@Operations.implementation_for(CreateBM25IndexOp)
+def _create_bm25_index_impl(operations: Operations, operation: CreateBM25IndexOp) -> None:
+    fields_sql = ", ".join(_quote_ident(field) for field in operation.fields)
+    sql = (
+        f"CREATE INDEX {_quote_ident(operation.index_name)} ON {_quote_ident(operation.table_name)} "
+        f"USING bm25 ({fields_sql}) WITH (key_field='{operation.key_field}')"
+    )
+    operations.execute(sql)
+
+
+@Operations.register_operation("drop_bm25_index")
+class DropBM25IndexOp(MigrateOperation):
+    def __init__(self, index_name: str, if_exists: bool = True) -> None:
+        self.index_name = index_name
+        self.if_exists = if_exists
+
+    @classmethod
+    def drop_bm25_index(cls, operations: Operations, index_name: str, if_exists: bool = True) -> MigrateOperation:
+        return operations.invoke(cls(index_name=index_name, if_exists=if_exists))
+
+
+@Operations.implementation_for(DropBM25IndexOp)
+def _drop_bm25_index_impl(operations: Operations, operation: DropBM25IndexOp) -> None:
+    if_exists_sql = " IF EXISTS" if operation.if_exists else ""
+    operations.execute(f"DROP INDEX{if_exists_sql} {_quote_ident(operation.index_name)}")
+
+
+@Operations.register_operation("reindex_bm25")
+class ReindexBM25Op(MigrateOperation):
+    def __init__(self, index_name: str, concurrently: bool = False) -> None:
+        self.index_name = index_name
+        self.concurrently = concurrently
+
+    @classmethod
+    def reindex_bm25(cls, operations: Operations, index_name: str, concurrently: bool = False) -> MigrateOperation:
+        return operations.invoke(cls(index_name=index_name, concurrently=concurrently))
+
+
+@Operations.implementation_for(ReindexBM25Op)
+def _reindex_bm25_impl(operations: Operations, operation: ReindexBM25Op) -> None:
+    concurrently_sql = " CONCURRENTLY" if operation.concurrently else ""
+    operations.execute(f"REINDEX INDEX{concurrently_sql} {_quote_ident(operation.index_name)}")

paradedb/sqlalchemy/errors.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+
+class ParadeDBError(Exception):
+    """Base class for ParadeDB SQLAlchemy integration errors."""
+
+
+class InvalidArgumentError(ParadeDBError, ValueError):
+    """Raised when a helper receives invalid user arguments."""
+
+
+class BM25ValidationError(ParadeDBError, ValueError):
+    """Base class for BM25 index validation errors."""
+
+
+class MissingKeyFieldError(BM25ValidationError):
+    """Raised when a BM25 index is missing key_field option."""
+
+
+class InvalidKeyFieldError(BM25ValidationError):
+    """Raised when BM25 key_field is not part of index fields."""
+
+
+class DuplicateTokenizerAliasError(BM25ValidationError):
+    """Raised when tokenizer aliases are duplicated in one BM25 index."""
+
+
+class InvalidBM25FieldError(BM25ValidationError):
+    """Raised when non-BM25Field expressions are used in a BM25 index."""
+
+
+class RuntimeGuardError(ParadeDBError, ValueError):
+    """Base class for runtime guardrail violations on statement builders."""
+
+
+class SnippetWithFuzzyPredicateError(RuntimeGuardError):
+    """Raised when snippet/snippets helpers are used with fuzzy predicates."""

paradedb/sqlalchemy/expr.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from sqlalchemy import func, literal
+from sqlalchemy.sql.elements import ClauseElement
+
+from .indexing import TokenizerSpec
+
+
+def json_text(json_expr: ClauseElement, key: str) -> ClauseElement:
+    return json_expr.op("->>")(literal(key))
+
+
+def concat_ws(separator: str, *parts: ClauseElement) -> ClauseElement:
+    return func.concat_ws(separator, *parts)
+
+
+def tokenizer_alias(tokenizer: TokenizerSpec | None) -> str | None:
+    return None if tokenizer is None else tokenizer.alias

paradedb/sqlalchemy/indexing.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from dataclasses import dataclass
 from typing import Any
 
@@ -10,6 +11,14 @@
 from sqlalchemy.sql.elements import ClauseElement, ColumnElement
 from sqlalchemy.sql.visitors import InternalTraversal
 
+from .errors import (
+    DuplicateTokenizerAliasError,
+    InvalidArgumentError,
+    InvalidBM25FieldError,
+    InvalidKeyFieldError,
+    MissingKeyFieldError,
+)
+
 
 @dataclass(frozen=True)
 class TokenizerSpec:
@@ -23,7 +32,7 @@ def render(self) -> str:
             return self.raw_sql
 
         if self.name is None:
-            raise ValueError("tokenizer name is required unless raw_sql is provided")
+            raise InvalidArgumentError("tokenizer name is required unless raw_sql is provided")
 
         if not self.options:
             return f"pdb.{self.name}()"
@@ -135,28 +144,28 @@ def validate_bm25_index(index: Index) -> None:
         return
 
     if not index.expressions:
-        raise ValueError("BM25 indexes must include at least one BM25Field")
+        raise InvalidBM25FieldError("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")
+        raise InvalidBM25FieldError("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")
+            raise DuplicateTokenizerAliasError(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>'}")
+        raise MissingKeyFieldError("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")
+        raise InvalidKeyFieldError(f"BM25 key_field '{key_field}' must match one of the indexed BM25Field columns")
 
 
 @event.listens_for(Index, "before_create")
@@ -172,6 +181,99 @@ class IndexMeta:
     aliases: dict[str, str]
 
 
+_KEY_FIELD_RE = re.compile(r"key_field\s*=\s*'?\"?([^'\",)\s]+)\"?'?", re.IGNORECASE)
+_ALIAS_RE = re.compile(r"alias\s*=\s*([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE)
+_CAST_FIELD_RE = re.compile(r"^\(*\"?([A-Za-z_][A-Za-z0-9_]*)\"?\)*\s*::\s*pdb\.", re.IGNORECASE)
+_PLAIN_FIELD_RE = re.compile(r'^\(*"?([A-Za-z_][A-Za-z0-9_]*)"?\)*$')
+
+
+def _split_top_level_csv(expr: str) -> list[str]:
+    parts: list[str] = []
+    current: list[str] = []
+    depth = 0
+    in_single = False
+    in_double = False
+
+    for ch in expr:
+        if ch == "'" and not in_double:
+            in_single = not in_single
+            current.append(ch)
+            continue
+        if ch == '"' and not in_single:
+            in_double = not in_double
+            current.append(ch)
+            continue
+        if not in_single and not in_double:
+            if ch == "(":
+                depth += 1
+            elif ch == ")":
+                depth = max(0, depth - 1)
+            elif ch == "," and depth == 0:
+                piece = "".join(current).strip()
+                if piece:
+                    parts.append(piece)
+                current = []
+                continue
+        current.append(ch)
+
+    tail = "".join(current).strip()
+    if tail:
+        parts.append(tail)
+    return parts
+
+
+def _extract_bm25_field_list(indexdef: str) -> list[str]:
+    marker = re.search(r"USING\s+bm25\s*\(", indexdef, re.IGNORECASE)
+    if marker is None:
+        return []
+
+    start = marker.end()
+    depth = 1
+    in_single = False
+    in_double = False
+    i = start
+    while i < len(indexdef):
+        ch = indexdef[i]
+        if ch == "'" and not in_double:
+            in_single = not in_single
+        elif ch == '"' and not in_single:
+            in_double = not in_double
+        elif not in_single and not in_double:
+            if ch == "(":
+                depth += 1
+            elif ch == ")":
+                depth -= 1
+                if depth == 0:
+                    return _split_top_level_csv(indexdef[start:i])
+        i += 1
+    return []
+
+
+def _extract_field_name(field_expr: str) -> str | None:
+    expr = field_expr.strip()
+    cast_match = _CAST_FIELD_RE.match(expr)
+    if cast_match:
+        return cast_match.group(1)
+    plain_match = _PLAIN_FIELD_RE.match(expr)
+    if plain_match:
+        return plain_match.group(1)
+    return None
+
+
+def _extract_key_field(indexdef: str) -> str | None:
+    match = _KEY_FIELD_RE.search(indexdef)
+    if match:
+        return match.group(1)
+    return None
+
+
+def _extract_alias(index_expr: str) -> str | None:
+    match = _ALIAS_RE.search(index_expr)
+    if match:
+        return match.group(1)
+    return None
+
+
 def describe(engine: Engine, table) -> list[IndexMeta]:
     query = text(
         """
@@ -188,21 +290,26 @@ def describe(engine: Engine, table) -> list[IndexMeta]:
     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]
+        key_field = _extract_key_field(indexdef)
+        raw_fields = _extract_bm25_field_list(indexdef)
+        aliases: dict[str, str] = {}
+        fields_ordered: list[str] = []
+        for raw in raw_fields:
+            field_name = _extract_field_name(raw)
+            if field_name is None:
+                continue
+            if field_name not in fields_ordered:
+                fields_ordered.append(field_name)
+            alias = _extract_alias(raw)
+            if alias is not None:
+                aliases[alias] = field_name
 
         output.append(
             IndexMeta(
                 index_name=row.indexname,
                 key_field=key_field,
-                fields=(),
-                aliases={},
+                fields=tuple(fields_ordered),
+                aliases=aliases,
             )
         )
     return output

paradedb/sqlalchemy/inspect.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Any
+
+from sqlalchemy.sql import visitors
+
+from ._pdb_cast import PDBCast
+
+_PARADEDB_PREDICATE_OPS = {"|||", "&&&", "###", "@@@", "==="}
+
+
+def collect_paradedb_operators(clause: Any) -> set[str]:
+    operators: set[str] = set()
+
+    def visit_binary(binary) -> None:
+        opstring = getattr(binary.operator, "opstring", None)
+        if opstring in _PARADEDB_PREDICATE_OPS:
+            operators.add(opstring)
+
+    visitors.traverse(clause, {}, {"binary": visit_binary})
+    return operators
+
+
+def has_paradedb_predicate(clause: Any) -> bool:
+    return bool(collect_paradedb_operators(clause))
+
+
+def _contains_fuzzy_cast(expr: Any) -> bool:
+    if isinstance(expr, PDBCast):
+        if expr.type_name == "fuzzy":
+            return True
+        return _contains_fuzzy_cast(expr.expr)
+    return False
+
+
+def has_fuzzy_predicate(clause: Any) -> bool:
+    found = False
+
+    def visit_binary(binary) -> None:
+        nonlocal found
+        if found:
+            return
+        opstring = getattr(binary.operator, "opstring", None)
+        if opstring != "===":
+            return
+        right = getattr(binary, "right", None)
+        if _contains_fuzzy_cast(right):
+            found = True
+
+    visitors.traverse(clause, {}, {"binary": visit_binary})
+    return found