From e753137a0da06d1db0c8b8d28580b54fef89c9f4 Mon Sep 17 00:00:00 2001 From: Ankit <573824+ankitml@users.noreply.github.com> Date: Wed, 4 Mar 2026 18:01:17 -0500 Subject: [PATCH] Add final DSL support so no rawsql needed for docs --- paradedb/sqlalchemy/search.py | 126 ++++++++++-- .../test_paradedb_queries_integration.py | 181 ++++++++++++++++++ tests/unit/test_sql_compilation_unit.py | 58 ++++++ tests/unit/test_validation_cache_unit.py | 15 ++ 4 files changed, 364 insertions(+), 16 deletions(-) diff --git a/paradedb/sqlalchemy/search.py b/paradedb/sqlalchemy/search.py index acd71bc..98d354f 100644 --- a/paradedb/sqlalchemy/search.py +++ b/paradedb/sqlalchemy/search.py @@ -1,10 +1,12 @@ from __future__ import annotations import json +import re +from collections.abc import Sequence from typing import Any -from sqlalchemy import Text, func, literal, literal_column, or_ -from sqlalchemy.dialects.postgresql import array +from sqlalchemy import Text, cast, func, literal, literal_column, or_ +from sqlalchemy.dialects.postgresql import ARRAY, array from sqlalchemy.sql import operators from sqlalchemy.sql.elements import ClauseElement, ColumnElement @@ -32,6 +34,7 @@ _QUERY: Any = operators.custom_op("@@@", precedence=5, is_comparison=True) _NEAR: Any = operators.custom_op("##", precedence=5) _NEAR_ORDERED: Any = operators.custom_op("##>", precedence=5) +_PDB_IDENTIFIER_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") def _to_term_payload(*terms: str) -> ClauseElement: @@ -49,6 +52,54 @@ def _apply_boost(expr: ClauseElement, boost: float | None) -> ClauseElement: return PDBCast(expr, "boost", (boost,)) +def _apply_const(expr: ClauseElement, const: float | None) -> ClauseElement: + if const is None: + return expr + if isinstance(expr, PDBCast) and expr.type_name in {"fuzzy", "slop"}: + expr = PDBCast(expr, "query") + return PDBCast(expr, "const", (const,)) + + +def _validate_pdb_identifier(name: str, *, field_name: str) -> str: + require_non_empty_string(name, field_name=field_name) + if not _PDB_IDENTIFIER_RE.fullmatch(name): + raise InvalidArgumentError( + f"{field_name} must be a bare identifier (letters, digits, underscore) for safe pdb cast rendering" + ) + return name + + +def _apply_tokenizer(expr: ClauseElement, tokenizer: str | None) -> ClauseElement: + if tokenizer is None: + return expr + tokenizer_name = _validate_pdb_identifier(tokenizer, field_name="tokenizer") + return PDBCast(expr, tokenizer_name) + + +def _to_phrase_payload(value: str | Sequence[str]) -> ClauseElement: + if isinstance(value, str): + require_non_empty_string(value, field_name="value") + return literal(value) + if not isinstance(value, Sequence): + raise InvalidArgumentError("value must be a non-empty string or a sequence of non-empty strings") + if not value: + raise InvalidArgumentError("value must contain at least one token") + require_non_empty_strings(value, field_name="value") + return array(list(value), type_=Text()) + + +def _apply_score_tuning( + expr: ClauseElement, + *, + boost: float | None = None, + const: float | None = None, +) -> ClauseElement: + if boost is not None and const is not None: + raise InvalidArgumentError("boost and const are mutually exclusive") + expr = _apply_boost(expr, boost) + return _apply_const(expr, const) + + def _apply_fuzzy( expr: ClauseElement, *, @@ -74,13 +125,16 @@ def match_all( field: ColumnElement, *terms: str, boost: float | None = None, + const: float | None = None, distance: int | None = None, prefix: bool = False, transpose_cost_one: bool = False, + tokenizer: str | None = None, ) -> ColumnElement[bool]: payload = _to_term_payload(*terms) payload = _apply_fuzzy(payload, distance=distance, prefix=prefix, transpose_cost_one=transpose_cost_one) - payload = _apply_boost(payload, boost) + payload = _apply_tokenizer(payload, tokenizer) + payload = _apply_score_tuning(payload, boost=boost, const=const) return field.operate(_MATCH_ALL, payload) @@ -88,13 +142,16 @@ def match_any( field: ColumnElement, *terms: str, boost: float | None = None, + const: float | None = None, distance: int | None = None, prefix: bool = False, transpose_cost_one: bool = False, + tokenizer: str | None = None, ) -> ColumnElement[bool]: payload = _to_term_payload(*terms) payload = _apply_fuzzy(payload, distance=distance, prefix=prefix, transpose_cost_one=transpose_cost_one) - payload = _apply_boost(payload, boost) + payload = _apply_tokenizer(payload, tokenizer) + payload = _apply_score_tuning(payload, boost=boost, const=const) return field.operate(_MATCH_ANY, payload) @@ -102,40 +159,65 @@ def term( field: ColumnElement, value: str, boost: float | None = None, + const: float | None = None, *, distance: int | None = None, prefix: bool = False, transpose_cost_one: bool = False, + tokenizer: str | None = None, ) -> ColumnElement[bool]: require_non_empty_string(value, field_name="value") payload: ClauseElement = literal(value) payload = _apply_fuzzy(payload, distance=distance, prefix=prefix, transpose_cost_one=transpose_cost_one) - payload = _apply_boost(payload, boost) + payload = _apply_tokenizer(payload, tokenizer) + payload = _apply_score_tuning(payload, boost=boost, const=const) return field.operate(_TERM, payload) def phrase( - field: ColumnElement, value: str, *, slop: int | None = None, boost: float | None = None + field: ColumnElement, + value: str | Sequence[str], + *, + slop: int | None = None, + boost: float | None = None, + const: float | None = None, + tokenizer: str | None = None, ) -> ColumnElement[bool]: - require_non_empty_string(value, field_name="value") if slop is not None: require_non_negative(slop, field_name="slop") - payload: ClauseElement = literal(value) + is_token_array = not isinstance(value, str) + payload: ClauseElement = _to_phrase_payload(value) + payload = _apply_tokenizer(payload, tokenizer) if slop is not None: + # psycopg binds array elements as VARCHAR by default; slop cast requires TEXT[]. + if is_token_array: + payload = cast(payload, ARRAY(Text())) payload = PDBCast(payload, "slop", (slop,)) - payload = _apply_boost(payload, boost) + payload = _apply_score_tuning(payload, boost=boost, const=const) return field.operate(_PHRASE, payload) -def regex(field: ColumnElement, pattern: str) -> ColumnElement[bool]: +def regex( + field: ColumnElement, + pattern: str, + *, + boost: float | None = None, + const: float | None = None, +) -> ColumnElement[bool]: require_non_empty_string(pattern, field_name="pattern") - return field.operate(_QUERY, func.pdb.regex(pattern)) + payload: ClauseElement = func.pdb.regex(pattern) + payload = _apply_score_tuning(payload, boost=boost, const=const) + return field.operate(_QUERY, payload) def all(field: ColumnElement) -> ColumnElement[bool]: return field.operate(_QUERY, func.pdb.all()) +def exists(field: ColumnElement) -> ColumnElement[bool]: + return field.operate(_QUERY, func.pdb.exists()) + + class ProximityExpr: def __init__(self, expr: ClauseElement) -> None: self.expr = expr @@ -264,7 +346,7 @@ def proximity(field: ColumnElement, prox: ProximityExpr | ClauseElement) -> Colu def range_term( field: ColumnElement, - bounds: str, + bounds: str | ClauseElement | int | float, *, relation: str = "Intersects", range_type: str | None = None, @@ -273,19 +355,31 @@ def range_term( Args: field: A range-typed column (int4range, daterange, tstzrange, etc.). - bounds: A range literal string, e.g. ``"[3,9]"``, ``"(3,9]"``. + bounds: Either a scalar point (e.g. ``1``) or a range literal string + (e.g. ``"[3,9]"``, ``"(3,9]"``). relation: One of ``"Intersects"``, ``"Contains"``, ``"Within"``, ``"ContainsOrIntersects"``. Defaults to ``"Intersects"``. + Only used when ``bounds`` is a range literal string. range_type: Optional PostgreSQL range type for explicit casting, e.g. ``"int4range"``, ``"int8range"``, ``"numrange"``, ``"daterange"``, ``"tsrange"``, ``"tstzrange"``. When provided, generates ``'bounds'::range_type`` cast. + Only used when ``bounds`` is a range literal string. Generates:: + field @@@ pdb.range_term(1) field @@@ pdb.range_term('[3,9]', 'Contains') field @@@ pdb.range_term('[3,9]'::int4range, 'Contains') """ + if not isinstance(bounds, str): + if relation != "Intersects": + raise InvalidArgumentError("relation is only supported when bounds is a range literal string") + if range_type is not None: + raise InvalidArgumentError("range_type is only supported when bounds is a range literal string") + scalar_arg = bounds if isinstance(bounds, ClauseElement) else literal(bounds) + return field.operate(_QUERY, func.pdb.range_term(scalar_arg)) + require_non_empty_string(bounds, field_name="bounds") if relation not in _VALID_RANGE_RELATIONS: raise InvalidArgumentError(f"relation must be one of: {', '.join(sorted(_VALID_RANGE_RELATIONS))}") @@ -295,10 +389,10 @@ def range_term( if range_type not in _VALID_RANGE_TYPES: raise InvalidArgumentError(f"range_type must be one of: {', '.join(sorted(_VALID_RANGE_TYPES))}") escaped = bounds.replace("'", "''") - bounds_arg: ClauseElement = literal_column(f"'{escaped}'::{range_type}") + range_bounds_arg: ClauseElement = literal_column(f"'{escaped}'::{range_type}") else: - bounds_arg = literal(bounds) - return field.operate(_QUERY, func.pdb.range_term(bounds_arg, relation_arg)) + range_bounds_arg = literal(bounds) + return field.operate(_QUERY, func.pdb.range_term(range_bounds_arg, relation_arg)) def more_like_this( diff --git a/tests/integration/test_paradedb_queries_integration.py b/tests/integration/test_paradedb_queries_integration.py index 8102763..5f608dc 100644 --- a/tests/integration/test_paradedb_queries_integration.py +++ b/tests/integration/test_paradedb_queries_integration.py @@ -42,6 +42,19 @@ def test_match_any_or_semantics(mock_session): assert ids == RUNNING_OR_WIRELESS_IDS +def test_match_any_custom_tokenizer(mock_session): + """match_any(..., tokenizer=) uses explicit query tokenization.""" + stmt_default = select(MockItem.id).where(search.match_any(MockItem.description, "running shoes")) + stmt_custom = select(MockItem.id).where( + search.match_any(MockItem.description, "running shoes", tokenizer="whitespace") + ) + assert_uses_paradedb_scan(mock_session, stmt_custom, index_name="mock_items_bm25_idx") + ids_default = _ids(mock_session, stmt_default) + ids_custom = _ids(mock_session, stmt_custom) + assert ids_default == SHOES_IDS + assert ids_custom == SHOES_IDS + + def test_match_all_and_semantics(mock_session): """match_all requires all terms to be present (AND search).""" stmt_all = select(MockItem.id).where(search.match_all(MockItem.description, "running", "shoes")) @@ -80,6 +93,48 @@ def test_phrase_with_slop(mock_session): assert ids_slop == RUNNING_IDS +def test_phrase_with_slop_and_const(mock_session): + """phrase(..., slop=, const=) bridges through query and executes.""" + stmt = ( + select(MockItem.id, pdb.score(MockItem.id).label("score")) + .where(search.phrase(MockItem.description, "running shoes", slop=2, const=1.0)) + .order_by(MockItem.id) + ) + assert_uses_paradedb_scan(mock_session, stmt, index_name="mock_items_bm25_idx") + rows = mock_session.execute(stmt).all() + assert [row.id for row in rows] == sorted(RUNNING_IDS) + assert [row.score for row in rows] == [pytest.approx(1.0)] + + +def test_phrase_custom_tokenizer(mock_session): + """phrase(..., tokenizer=) uses explicit query tokenization.""" + stmt_default = select(MockItem.id).where(search.phrase(MockItem.description, "running shoes")) + stmt_custom = select(MockItem.id).where( + search.phrase(MockItem.description, "running shoes", tokenizer="whitespace") + ) + assert_uses_paradedb_scan(mock_session, stmt_custom, index_name="mock_items_bm25_idx") + ids_default = _ids(mock_session, stmt_default) + ids_custom = _ids(mock_session, stmt_custom) + assert ids_default == RUNNING_IDS + assert ids_custom == RUNNING_IDS + + +def test_phrase_pretokenized(mock_session): + """phrase() accepts explicit token arrays.""" + stmt = select(MockItem.id).where(search.phrase(MockItem.description, ["running", "shoes"])) + assert_uses_paradedb_scan(mock_session, stmt, index_name="mock_items_bm25_idx") + ids = _ids(mock_session, stmt) + assert ids == RUNNING_IDS + + +def test_phrase_pretokenized_with_slop(mock_session): + """phrase() supports slop for token arrays.""" + stmt = select(MockItem.id).where(search.phrase(MockItem.description, ["shoes", "running"], slop=2)) + assert_uses_paradedb_scan(mock_session, stmt, index_name="mock_items_bm25_idx") + ids = _ids(mock_session, stmt) + assert ids == RUNNING_IDS + + def test_regex_match(mock_session): """regex() matches patterns against indexed tokens.""" regex_stmt = select(MockItem.id).where(search.regex(MockItem.description, "run.*")) @@ -88,6 +143,17 @@ def test_regex_match(mock_session): assert regex_ids == RUNNING_IDS +def test_regex_boost_preserves_matches(mock_session): + """regex(..., boost=) only changes scoring, not match set.""" + stmt_base = select(MockItem.id).where(search.regex(MockItem.description, "run.*")) + stmt_boost = select(MockItem.id).where(search.regex(MockItem.description, "run.*", boost=2.0)) + assert_uses_paradedb_scan(mock_session, stmt_boost, index_name="mock_items_bm25_idx") + ids_base = _ids(mock_session, stmt_base) + ids_boost = _ids(mock_session, stmt_boost) + assert ids_base == RUNNING_IDS + assert ids_boost == RUNNING_IDS + + # --------------------------------------------------------------------------- # Fuzzy search # --------------------------------------------------------------------------- @@ -140,6 +206,32 @@ def test_fuzzy_with_boost(mock_session): assert ids_base == ids_boost +def test_fuzzy_with_const(mock_session): + """match_any(..., distance=, const=) bridges through query and executes.""" + stmt = ( + select(MockItem.id, pdb.score(MockItem.id).label("score")) + .where(search.match_any(MockItem.description, "runnning", distance=1, const=1.0)) + .order_by(MockItem.id) + ) + assert_uses_paradedb_scan(mock_session, stmt, index_name="mock_items_bm25_idx") + rows = mock_session.execute(stmt).all() + assert [row.id for row in rows] == sorted(RUNNING_IDS) + assert [row.score for row in rows] == [pytest.approx(1.0)] + + +def test_match_any_const_sets_constant_score(mock_session): + """match_any(..., const=) assigns a constant score to every match.""" + stmt = ( + select(MockItem.id, pdb.score(MockItem.id).label("score")) + .where(search.match_any(MockItem.description, "shoes", const=1.0)) + .order_by(MockItem.id) + ) + assert_uses_paradedb_scan(mock_session, stmt, index_name="mock_items_bm25_idx") + rows = mock_session.execute(stmt).all() + assert [row.id for row in rows] == sorted(SHOES_IDS) + assert [row.score for row in rows] == [pytest.approx(1.0), pytest.approx(1.0), pytest.approx(1.0)] + + # --------------------------------------------------------------------------- # Parse / query-string search # --------------------------------------------------------------------------- @@ -460,6 +552,54 @@ def test_range_term_with_range_type(mock_session): conn.execute(text("DROP TABLE IF EXISTS rt_items_pq")) +def test_range_term_scalar_contains_point(mock_session): + """range_term() supports scalar point queries like pdb.range_term(1).""" + from sqlalchemy import Column, Integer, MetaData, Table, text + from sqlalchemy.dialects.postgresql import INT4RANGE + + engine = mock_session.get_bind() + metadata = MetaData() + tbl = Table( + "rt_scalar_items_pq", + metadata, + Column("id", Integer, primary_key=True), + Column("weight_range", INT4RANGE, nullable=False), + ) + + with engine.begin() as conn: + conn.execute(text("DROP INDEX IF EXISTS rt_scalar_items_pq_bm25_idx")) + conn.execute(text("DROP TABLE IF EXISTS rt_scalar_items_pq")) + metadata.create_all(engine) + + with engine.begin() as conn: + conn.execute( + text( + "CREATE INDEX rt_scalar_items_pq_bm25_idx ON rt_scalar_items_pq USING bm25 (id, weight_range) " + "WITH (key_field='id')" + ) + ) + conn.execute( + text( + "INSERT INTO rt_scalar_items_pq (id, weight_range) VALUES " + "(1, '[1,4]'::int4range), (2, '[3,9]'::int4range), (3, '[10,12]'::int4range)" + ) + ) + + try: + with engine.connect() as conn: + from sqlalchemy.orm import Session as S + + with S(bind=conn) as s: + stmt = select(tbl.c.id).where(search.range_term(tbl.c.weight_range, 1)).order_by(tbl.c.id) + assert_uses_paradedb_scan(s, stmt, index_name="rt_scalar_items_pq_bm25_idx") + ids = list(s.scalars(stmt)) + assert ids == [1] + finally: + with engine.begin() as conn: + conn.execute(text("DROP INDEX IF EXISTS rt_scalar_items_pq_bm25_idx")) + conn.execute(text("DROP TABLE IF EXISTS rt_scalar_items_pq")) + + def test_range_term_invalid_range_type_raises(): """range_term() with unknown range_type raises InvalidArgumentError.""" from sqlalchemy import Column, Integer, MetaData, Table @@ -508,3 +648,44 @@ def test_all_predicate_matches_everything(mock_session): select(MockItem.id).where(search.all(MockItem.id)).with_only_columns(__import__("sqlalchemy").func.count()) ) assert total == ALL_MOCK_ITEM_COUNT + + +def test_exists_query_matches_non_null(mock_session): + """search.exists() returns only documents where the indexed field is non-null.""" + from sqlalchemy import Column, Integer, MetaData, Table, text + + engine = mock_session.get_bind() + metadata = MetaData() + tbl = Table( + "exists_items_pq", + metadata, + Column("id", Integer, primary_key=True), + Column("rating", Integer, nullable=True), + ) + + with engine.begin() as conn: + conn.execute(text("DROP INDEX IF EXISTS exists_items_pq_bm25_idx")) + conn.execute(text("DROP TABLE IF EXISTS exists_items_pq")) + metadata.create_all(engine) + + with engine.begin() as conn: + conn.execute( + text( + "CREATE INDEX exists_items_pq_bm25_idx ON exists_items_pq USING bm25 (id, rating) WITH (key_field='id')" + ) + ) + conn.execute(text("INSERT INTO exists_items_pq (id, rating) VALUES (1, 5), (2, NULL), (3, 0)")) + + try: + with engine.connect() as conn: + from sqlalchemy.orm import Session as S + + with S(bind=conn) as s: + stmt = select(tbl.c.id).where(search.exists(tbl.c.rating)).order_by(tbl.c.id) + assert_uses_paradedb_scan(s, stmt, index_name="exists_items_pq_bm25_idx") + ids = list(s.scalars(stmt)) + assert ids == [1, 3] + finally: + with engine.begin() as conn: + conn.execute(text("DROP INDEX IF EXISTS exists_items_pq_bm25_idx")) + conn.execute(text("DROP TABLE IF EXISTS exists_items_pq")) diff --git a/tests/unit/test_sql_compilation_unit.py b/tests/unit/test_sql_compilation_unit.py index b9bead6..974ea6b 100644 --- a/tests/unit/test_sql_compilation_unit.py +++ b/tests/unit/test_sql_compilation_unit.py @@ -14,6 +14,7 @@ column("id", Integer), column("description", Text), column("category", String), + column("rating", Integer), ) @@ -73,6 +74,56 @@ def test_regex_and_all_compile(): assert "id @@@ pdb.all()" in all_sql +def test_match_any_with_tokenizer_compile(): + stmt = select(products.c.id).where( + search.match_any(products.c.description, "running shoes", tokenizer="whitespace") + ) + sql = _sql(stmt) + assert "description ||| 'running shoes'::pdb.whitespace" in sql + + +def test_phrase_with_tokenizer_compile(): + stmt = select(products.c.id).where(search.phrase(products.c.description, "running shoes", tokenizer="whitespace")) + sql = _sql(stmt) + assert "description ### 'running shoes'::pdb.whitespace" in sql + + +def test_phrase_pretokenized_with_slop_compile(): + stmt = select(products.c.id).where(search.phrase(products.c.description, ["shoes", "running"], slop=2)) + sql = _sql(stmt) + assert "description ### CAST(ARRAY['shoes', 'running'] AS TEXT[])::pdb.slop(2)" in sql + + +def test_phrase_with_slop_and_const_compile(): + stmt = select(products.c.id).where(search.phrase(products.c.description, "running shoes", slop=2, const=1.0)) + sql = _sql(stmt) + assert "description ### 'running shoes'::pdb.slop(2)::pdb.query::pdb.const(1.0)" in sql + + +def test_regex_boost_compile(): + stmt = select(products.c.id).where(search.regex(products.c.description, "key.*", boost=2.0)) + sql = _sql(stmt) + assert "description @@@ pdb.regex('key.*')::pdb.boost(2.0)" in sql + + +def test_match_any_fuzzy_with_const_compile(): + stmt = select(products.c.id).where(search.match_any(products.c.description, "shose", distance=2, const=1.0)) + sql = _sql(stmt) + assert "description ||| 'shose'::pdb.fuzzy(2)::pdb.query::pdb.const(1.0)" in sql + + +def test_match_any_const_compile(): + stmt = select(products.c.id).where(search.match_any(products.c.description, "shoes", const=1.0)) + sql = _sql(stmt) + assert "description ||| 'shoes'::pdb.const(1.0)" in sql + + +def test_exists_compile(): + stmt = select(products.c.id).where(search.exists(products.c.rating)) + sql = _sql(stmt) + assert "rating @@@ pdb.exists()" in sql + + def test_pdb_helpers_compile(): stmt = select( pdb.score(products.c.id).label("score"), @@ -262,6 +313,13 @@ def test_range_term_invalid_relation_raises(): search.range_term(products.c.id, "[3,9]", relation="BadRelation") +def test_range_term_scalar_compile(): + range_items = table("range_items", column("id", Integer), column("weight_range", postgresql.INT4RANGE)) + stmt = select(range_items.c.id).where(search.range_term(range_items.c.weight_range, 1)) + sql = _sql(stmt) + assert "weight_range @@@ pdb.range_term(1)" in sql + + # --------------------------------------------------------------------------- # New feature: indexing.validate_pushdown # --------------------------------------------------------------------------- diff --git a/tests/unit/test_validation_cache_unit.py b/tests/unit/test_validation_cache_unit.py index f1c4f80..5f61976 100644 --- a/tests/unit/test_validation_cache_unit.py +++ b/tests/unit/test_validation_cache_unit.py @@ -72,6 +72,21 @@ def test_search_argument_validation_errors(): with pytest.raises(InvalidArgumentError, match="pattern must be a non-empty string"): search.regex(products.c.description, "") + with pytest.raises(InvalidArgumentError, match="tokenizer must be a bare identifier"): + search.match_any(products.c.description, "running shoes", tokenizer="whitespace;drop") + + with pytest.raises(InvalidArgumentError, match="value must contain at least one token"): + search.phrase(products.c.description, []) + + with pytest.raises(InvalidArgumentError, match="range_type is only supported"): + search.range_term(products.c.id, 1, range_type="int4range") + + with pytest.raises(InvalidArgumentError, match="relation is only supported"): + search.range_term(products.c.id, 1, relation="Contains") + + with pytest.raises(InvalidArgumentError, match="mutually exclusive"): + search.match_any(products.c.description, "shoes", boost=2.0, const=1.0) + def test_more_like_this_uses_specific_error_type(): with pytest.raises(InvalidMoreLikeThisOptionsError, match="exactly one"):