diff --git a/DESIGN.md b/DESIGN.md index 982ec8d..09aaf0a 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -209,7 +209,8 @@ write-through for callers that need `.sql()` on the entry immediately. ## MaterializedFrame -Implemented today: `.arrow()`, `.df()`, `.to_pylist()`, `len()`, and +Implemented today: `.arrow()`, `.df()`, `.to_pylist()`, `len()`, +`.search()`, `.vector_search()`, and `.sql()` (server-side; requires a persisted entry — on a fresh write-behind miss it raises until the background persist lands). The rest of this section is the target design. @@ -253,6 +254,7 @@ from hotdata_materialized import materialize version=2, # bump to bust on code changes background=True, # write-behind (default) | write-through key_fn=None, # stable identity for opaque arguments + index=BM25("notes"), # or Vector(...); built at persist time ) def daily_signups(start, end): return (Event.objects @@ -261,11 +263,11 @@ def daily_signups(start, end): .annotate(n=Count("id"))) ``` -Planned knobs (steps 4–5, do not document user-facing until they exist): +Planned knobs (step 4, do not document user-facing until they exist): `mode="swr" | "strict"` (refresh behavior), `invalidate_on=[Model]` -(signal-based staleness), `index=Vector(...)/BM25(...)` (search indexes), -`on_error="fallback" | "raise"` (fail-open is currently always on), and a -`from_queryset(qs, ttl=)` imperative helper for one-off QuerySets. +(signal-based staleness), `on_error="fallback" | "raise"` (fail-open is +currently always on), and a `from_queryset(qs, ttl=)` imperative helper +for one-off QuerySets. ## Fingerprinting diff --git a/README.md b/README.md index 06d6018..63bb8fe 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,42 @@ Fail-open by design: if Hotdata is unreachable, the function runs and its result is served uncached — the cache degrades to "no cache," never to "no page." +## Search + +Declare a search index on the decorator and the entry gets it when it +persists (builds run server-side; a freshly missed entry may error on search +for a few seconds until its index is ready): + +```python +from hotdata_materialized import BM25, Vector, materialize + + +@materialize(ttl=3600, index=BM25("notes")) +def incidents(): + return Incident.objects.values("id", "notes", "severity") + + +incidents().search("checkout errors outage", column="notes", limit=5) +``` + +Vector (semantic) search embeds your text server-side via a workspace +embedding provider — pass the indexed source column and a natural-language +query: + +```python +@materialize(ttl=3600, index=Vector("notes", provider="emb_...", metric="cosine")) +def incidents_semantic(): + return Incident.objects.values("id", "notes") + + +incidents_semantic().vector_search("the site was down", column="notes", limit=5) +``` + +Both return the matching rows as a `pyarrow.Table` with a relevance column +(`score` / `distance`). One platform rule to know: an embedding-backed vector +index cannot share an entry with other indexes — the decorator rejects that +combination up front. + ## Evicting entries The decorator composes public primitives (`fingerprint_call` / @@ -142,7 +178,8 @@ trade-offs. - [ ] Stale-while-revalidate refresh (rebuild protocol) - [ ] Sweep command for expired entries - [ ] Chainable queryset facade on the frame (`.filter()/.order_by()`) -- [ ] Vector/BM25 index declarations and `frame.search()` +- [x] Vector/BM25 index declarations, `frame.search()` and + `frame.vector_search()` - [ ] Async view support (`await frame.aarrow()`) - [ ] PyPI release diff --git a/hotdata_materialized/__init__.py b/hotdata_materialized/__init__.py index 8c9329b..68ea32d 100644 --- a/hotdata_materialized/__init__.py +++ b/hotdata_materialized/__init__.py @@ -15,6 +15,7 @@ StoreError, ) from .fingerprint import fingerprint_call, fingerprint_queryset +from .indexes import BM25, Vector from .registry import Registry, RegistryEntry from .store import EntryStore @@ -23,6 +24,8 @@ __all__ = [ "materialize", "MaterializedFrame", + "BM25", + "Vector", "Config", "HotdataClients", "get_clients", diff --git a/hotdata_materialized/_sql.py b/hotdata_materialized/_sql.py index c12e56b..802d882 100644 --- a/hotdata_materialized/_sql.py +++ b/hotdata_materialized/_sql.py @@ -1,16 +1,26 @@ -"""SQL literal encoding for registry statements. +"""SQL literal and identifier encoding. POST /v1/queries takes SQL text only — there are no bind parameters — so -every value written into registry SQL must go through quote_literal(). +every value written into SQL must go through quote_literal(), and every +caller-supplied identifier (column names in search) through quote_ident(). """ from __future__ import annotations import datetime as _dt import math +import re from .exceptions import RegistryError +_IDENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z") + + +def quote_ident(name: str) -> str: + if not _IDENT.match(name): + raise ValueError(f"invalid SQL identifier: {name!r}") + return name + def quote_literal(value) -> str: if value is None: diff --git a/hotdata_materialized/client.py b/hotdata_materialized/client.py index 3e68ada..4323ae8 100644 --- a/hotdata_materialized/client.py +++ b/hotdata_materialized/client.py @@ -29,6 +29,7 @@ def __init__(self, config: Config): self.databases = hotdata.DatabasesApi(self.api_client) self.uploads = UploadsApi(self.api_client) self.results = ResultsApi(self.api_client) + self.indexes = hotdata.IndexesApi(self.api_client) _lock = threading.Lock() diff --git a/hotdata_materialized/decorator.py b/hotdata_materialized/decorator.py index 1c26b46..b243a30 100644 --- a/hotdata_materialized/decorator.py +++ b/hotdata_materialized/decorator.py @@ -17,16 +17,18 @@ import logging import re import threading -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union import pyarrow as pa from .client import get_clients from .conf import Config from .exceptions import MaterializedError +from ._sql import quote_ident, quote_literal from .fingerprint import fingerprint_call +from .indexes import BM25, Vector from .registry import STATUS_READY, Registry, RegistryEntry -from .store import DATA_TABLE, EntryStore +from .store import DATA_SCHEMA, DATA_TABLE, EntryStore logger = logging.getLogger(__name__) @@ -81,6 +83,29 @@ def sql(self, sql: str) -> pa.Table: then; use background=False on the decorator to persist inline).""" return self._store.query_table(self.entry, _rewrite_this(sql)) + def search(self, query: str, *, column: str, limit: int = 10) -> pa.Table: + """BM25 keyword search over the entry, best matches first. Requires + a BM25 index on the column (declare with @materialize(index=BM25(...)); + a freshly missed entry may error until the index build finishes).""" + table_ref = f"default.{DATA_SCHEMA}.{DATA_TABLE}" + return self._store.query_table( + self.entry, + f"SELECT * FROM bm25_search({quote_literal(table_ref)}, " + f"{quote_literal(quote_ident(column))}, {quote_literal(query)}) " + f"ORDER BY score DESC LIMIT {int(limit)}", + ) + + def vector_search(self, query: str, *, column: str, limit: int = 10) -> pa.Table: + """Semantic search over the entry, nearest first. Requires a vector + index (declare with @materialize(index=Vector(...))). Pass the indexed + source column; the query text is embedded server-side.""" + return self._store.query_table( + self.entry, + f"SELECT *, vector_distance({quote_ident(column)}, " + f"{quote_literal(query)}) AS distance " + f"FROM {DATA_TABLE} ORDER BY distance LIMIT {int(limit)}", + ) + def __len__(self) -> int: return self.arrow().num_rows @@ -131,6 +156,7 @@ def materialize( version: int = 0, background: Optional[bool] = None, key_fn: Optional[Callable[..., Any]] = None, + index: Union[BM25, Vector, Sequence[Union[BM25, Vector]], None] = None, ) -> Callable[[Callable[..., Any]], Callable[..., MaterializedFrame]]: """Materialize a function's result into Hotdata. @@ -140,8 +166,26 @@ def revenue_by_region(): The call returns a MaterializedFrame either way; `.cached` says which path served it. `version=` busts the cache on code changes; `key_fn=` maps - non-JSON-serializable arguments to a stable identity. + non-JSON-serializable arguments to a stable identity. `index=` declares + BM25/Vector search indexes built when the entry persists. """ + declarations: Tuple[Union[BM25, Vector], ...] + if index is None: + declarations = () + elif isinstance(index, (BM25, Vector)): + declarations = (index,) + else: + declarations = tuple(index) + embedded = any( + isinstance(d, Vector) and d.provider is not None for d in declarations + ) + if embedded and len(declarations) > 1: + # platform constraint: embedding-backed vector indexes cannot coexist + # with other indexes on the same table + raise ValueError( + "an embedding-backed Vector index (provider=...) cannot be " + "combined with other indexes on the same entry" + ) def decorate(func: Callable[..., Any]) -> Callable[..., MaterializedFrame]: @functools.wraps(func) @@ -177,6 +221,7 @@ def wrapper(*args: Any, **kwargs: Any) -> MaterializedFrame: ttl=ttl, version=version, background=background, + indexes=declarations, ) except MaterializedError: logger.warning( diff --git a/hotdata_materialized/indexes.py b/hotdata_materialized/indexes.py new file mode 100644 index 0000000..40177a3 --- /dev/null +++ b/hotdata_materialized/indexes.py @@ -0,0 +1,33 @@ +"""Search-index declarations for @materialize(index=...). + +Indexes are created right after the entry's data loads; builds run as +background jobs on the Hotdata side, so a freshly missed entry may briefly +error on search until its index is ready. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + + +@dataclass(frozen=True) +class BM25: + """Keyword (BM25) index over a text column; query via frame.search().""" + + column: str + + +@dataclass(frozen=True) +class Vector: + """Vector index; query via frame.vector_search(). + + With `provider` set (an embedding provider id), the column is treated + as text and embedded server-side; vector_search then targets the source + column and query text is embedded automatically. An embedding-backed + index cannot share a table with other indexes.""" + + column: str + provider: Optional[str] = None + metric: Optional[str] = None # "l2" | "cosine" | "dot" + dimensions: Optional[int] = None diff --git a/hotdata_materialized/store.py b/hotdata_materialized/store.py index ff73c70..beb7eb6 100644 --- a/hotdata_materialized/store.py +++ b/hotdata_materialized/store.py @@ -21,15 +21,17 @@ import datetime as _dt import logging import threading +import time from concurrent.futures import Future, ThreadPoolExecutor, wait as futures_wait -from typing import Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Sequence, Set, Union import pyarrow as pa from hotdata.arrow import ResultNotReadyError -from hotdata.exceptions import ApiException +from hotdata.exceptions import ApiException, ConflictException from hotdata.models.create_database_request import CreateDatabaseRequest from hotdata.models.database_default_schema_decl import DatabaseDefaultSchemaDecl from hotdata.models.database_default_table_decl import DatabaseDefaultTableDecl +from hotdata.models.create_index_request import CreateIndexRequest from hotdata.models.load_managed_table_request import LoadManagedTableRequest from hotdata.models.query_request import QueryRequest @@ -42,6 +44,7 @@ ) from .conf import Config from .exceptions import StoreError +from .indexes import BM25, Vector from .registry import STATUS_BUILDING, Registry, RegistryEntry, utcnow_iso __all__ = [ @@ -92,6 +95,7 @@ def materialize( ttl: Optional[int] = None, version: int = 0, background: Optional[bool] = None, + indexes: Optional[Sequence[Union[BM25, Vector]]] = None, ) -> RegistryEntry: """Persist a captured result as a materialized entry. @@ -101,12 +105,14 @@ def materialize( if background is None: background = self._config.background if not background: - return self._persist(fingerprint, table, key=key, ttl=ttl, version=version) + return self._persist( + fingerprint, table, key=key, ttl=ttl, version=version, indexes=indexes + ) executor = _get_executor(self._config.background_max_workers) with self._inflight_lock: future = executor.submit( - self._persist_background, fingerprint, table, key, ttl, version + self._persist_background, fingerprint, table, key, ttl, version, indexes ) self._inflight.setdefault(fingerprint, set()).add(future) @@ -137,9 +143,13 @@ def flush(self, timeout: Optional[float] = None) -> List[BaseException]: done, _ = futures_wait(pending, timeout=timeout) return [e for e in (f.exception() for f in done) if e is not None] - def _persist_background(self, fingerprint, table, key, ttl, version) -> None: + def _persist_background( + self, fingerprint, table, key, ttl, version, indexes + ) -> None: try: - self._persist(fingerprint, table, key=key, ttl=ttl, version=version) + self._persist( + fingerprint, table, key=key, ttl=ttl, version=version, indexes=indexes + ) except Exception: # Loud by design: a quietly failing write-behind path means the # cache never fills and nobody notices. @@ -157,9 +167,11 @@ def _persist( key: Optional[str] = None, ttl: Optional[int] = None, version: int = 0, + indexes: Optional[Sequence[Union[BM25, Vector]]] = None, ) -> RegistryEntry: parquet_bytes = table_to_parquet_bytes(table) - database_id = self._create_entry_database(fingerprint, ttl) + created = self._create_entry_database(fingerprint, ttl) + database_id = created.id try: upload = self._clients.uploads.upload_file( parquet_bytes, @@ -189,7 +201,7 @@ def _persist( x_database_id=database_id, auto_follow=False, ).result_id - return self._registry.mark_ready( + entry = self._registry.mark_ready( fingerprint, database_id=database_id, ttl=ttl, @@ -207,6 +219,20 @@ def _persist( self._cleanup_failed_build(fingerprint, database_id) raise StoreError(f"materializing entry failed: {exc}") from exc + # Index kickoff is deliberately OUTSIDE the publication failure + # domain: a bad provider id or a stuck table lock must not discard + # otherwise-valid cached data. The entry stays served; search errors + # at the search site until the index problem is fixed. + try: + for declaration in indexes or (): + self._create_index(created.default_connection_id, declaration) + except Exception as exc: + raise StoreError( + f"index creation for entry {fp.short(fingerprint)} failed " + f"(entry cached; search unavailable): {exc}" + ) from exc + return entry + def read_table(self, entry: RegistryEntry) -> pa.Table: """Materialize an entry's data as a pyarrow.Table. @@ -293,7 +319,7 @@ def evict(self, fingerprint: str) -> None: # -- internals ----------------------------------------------------------- - def _create_entry_database(self, fingerprint: str, ttl: Optional[int]) -> str: + def _create_entry_database(self, fingerprint: str, ttl: Optional[int]) -> Any: # Server-side expiry (ttl + grace) is a best-effort backstop behind # the sweep; label is display-only — identity is the id via registry. expires_at = None @@ -317,7 +343,41 @@ def _create_entry_database(self, fingerprint: str, ttl: Optional[int]) -> str: ) except Exception as exc: raise StoreError(f"creating entry database failed: {exc}") from exc - return created.id + return created + + def _create_index(self, connection_id: str, declaration: Union[BM25, Vector]) -> None: + # index builds run as async jobs server-side; the persist does not wait + # "async" is a reserved word, so the SDK model aliases it; build + # requests from the wire shape + if isinstance(declaration, BM25): + request = CreateIndexRequest.model_validate({ + "index_name": f"{DATA_TABLE}_{declaration.column}_bm25", + "index_type": "bm25", + "columns": [declaration.column], + "async": True, + }) + else: + request = CreateIndexRequest.model_validate({ + "index_name": f"{DATA_TABLE}_{declaration.column}_vector", + "index_type": "vector", + "columns": [declaration.column], + "metric": declaration.metric, + "embedding_provider_id": declaration.provider, + "dimensions": declaration.dimensions, + "async": True, + }) + # the data load can still hold the table lock when this runs + # (observed: 409 RESOURCE_LOCKED with Retry-After); wait it out + for attempt in range(6): + try: + self._clients.indexes.create_index( + connection_id, DATA_SCHEMA, DATA_TABLE, request + ) + return + except ConflictException: + if attempt == 5: + raise + time.sleep(5) def _delete_database(self, database_id: str) -> None: try: diff --git a/tests/fakes.py b/tests/fakes.py index d897279..b1853f4 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -49,6 +49,8 @@ def __init__(self): load_database_table=self._single_load, ) self.uploads = SimpleNamespace(upload_file=self._upload_file) + self.index_creates = [] # (connection_id, schema, table, request) + self.indexes = SimpleNamespace(create_index=self._create_index) # -- query --------------------------------------------------------------- @@ -104,6 +106,7 @@ def _create_database(self, request): name=request.name, default_catalog="default", default_schema="main", + default_connection_id=f"conn_{db_id}", ) def _delete_database(self, database_id): @@ -182,3 +185,8 @@ def _upload_file(self, source, *, filename=None, content_type=None, **kwargs): upload_id = f"up_{next(self._ids)}" self.upload_blobs[upload_id] = bytes(source) return SimpleNamespace(upload_id=upload_id, size_bytes=len(source)) + + def _create_index(self, connection_id, var_schema, table, request): + with self._api_lock: + self.index_creates.append((connection_id, var_schema, table, request)) + return SimpleNamespace(index_name=request.index_name, status="building") diff --git a/tests/test_decorator.py b/tests/test_decorator.py index c7edd11..cbb0510 100644 --- a/tests/test_decorator.py +++ b/tests/test_decorator.py @@ -146,3 +146,104 @@ def test_to_arrow_accepts_table_and_rows(): def test_to_arrow_rejects_non_dict_rows(): with pytest.raises(TypeError, match="iterable of dicts"): to_arrow([1, 2, 3]) + + +def test_index_declarations_create_indexes_at_persist(runtime): + from hotdata_materialized import BM25, Vector + + @materialize(ttl=60, index=BM25("notes")) + def notes(): + return [{"notes": "server outage in eu-west", "n": 1}] + + @materialize(ttl=60, index=Vector("notes", provider="emb_1", metric="cosine")) + def embedded_notes(): + return [{"notes": "payment gateway timeout", "n": 2}] + + notes() + embedded_notes() + runtime.store.flush() + assert len(runtime.backend.index_creates) == 2 + # two background persists race; order by type before asserting + (conn_id, schema, table, bm25_req), (_, _, _, vec_req) = sorted( + runtime.backend.index_creates, key=lambda c: c[3].index_type + ) + assert conn_id.startswith("conn_db_") + assert (schema, table) == ("main", "data") + assert bm25_req.index_type == "bm25" + assert bm25_req.columns == ["notes"] + assert bm25_req.var_async is True + assert vec_req.index_type == "vector" + assert vec_req.metric == "cosine" + assert vec_req.embedding_provider_id == "emb_1" + + +def test_no_index_declared_creates_none(runtime): + calls = {"n": 0} + report = make_report(calls) + report("2026-07-22") + runtime.store.flush() + assert runtime.backend.index_creates == [] + + +def test_search_builds_bm25_sql(runtime, monkeypatch): + calls = {"n": 0} + report = make_report(calls) + report("2026-07-22") + runtime.store.flush() + frame = report("2026-07-22") + captured: dict[str, str] = {} + monkeypatch.setattr( + runtime.store, "query_table", lambda entry, sql: captured.update(sql=sql) + ) + frame.search("outage in eu-west", column="day", limit=5) + assert captured["sql"] == ( + "SELECT * FROM bm25_search('default.main.data', 'day', " + "'outage in eu-west') ORDER BY score DESC LIMIT 5" + ) + + +def test_vector_search_builds_sql_and_validates_column(runtime, monkeypatch): + calls = {"n": 0} + report = make_report(calls) + report("2026-07-22") + runtime.store.flush() + frame = report("2026-07-22") + captured: dict[str, str] = {} + monkeypatch.setattr( + runtime.store, "query_table", lambda entry, sql: captured.update(sql=sql) + ) + frame.vector_search("churn risk", column="notes", limit=3) + assert captured["sql"] == ( + "SELECT *, vector_distance(notes, 'churn risk') AS distance " + "FROM data ORDER BY distance LIMIT 3" + ) + with pytest.raises(ValueError, match="identifier"): + frame.vector_search("x", column="notes; DROP TABLE data") + + +def test_embedded_vector_cannot_combine_with_other_indexes(): + from hotdata_materialized import BM25, Vector + + with pytest.raises(ValueError, match="cannot be combined"): + materialize(index=[BM25("notes"), Vector("notes", provider="emb_1")]) + + +def test_index_failure_does_not_discard_the_cached_entry(runtime, monkeypatch): + from hotdata_materialized import BM25 + from hotdata_materialized.registry import STATUS_READY + + def broken(*args, **kwargs): + raise RuntimeError("bad embedding provider") + + monkeypatch.setattr(runtime.backend.indexes, "create_index", broken) + + @materialize(ttl=60, index=BM25("notes")) + def notes(): + return [{"notes": "outage", "n": 1}] + + notes() + errors = runtime.store.flush() + assert len(errors) == 1 and "entry cached" in str(errors[0]) + frame = notes() # next call is a hit despite the failed index + assert frame.cached is True + assert frame.entry.status == STATUS_READY