diff --git a/CHANGELOG.md b/CHANGELOG.md index c318b0f..a6b3545 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Queries return their rows in one request when they finish within + `async_after_ms` (default 1000; `None` = always async), using the inline + Arrow response. Slower or larger queries still poll as before. +- `connect()` falls back to `HOTDATA_API_URL`, `HOTDATA_API_KEY` and + `HOTDATA_WORKSPACE` when `api_url`, `token` or `workspace_id` is not passed. + `database_id` and `async_after_ms` are also accepted in `hotdata://` URLs. +- `create_index`, `list_indexes` and `drop_index` on the backend, including + IVF vector indexes (`algorithm`, `nlist`, `probe_fraction`, + `vector_precision`) and embedding-provider indexes. +- `vector_distance` UDF, and `semantic_search` accepts a text query for + embedding-backed indexes. + +### Changed + +- Require `hotdata>=0.12,<0.13` (was `>=0.11,<0.12`). +- Instant-database table loads now retry while another load holds the table + lock (uses the SDK's top-level `ConnectionsApi`/`DatabasesApi`). + ## [0.6.0] - 2026-10-01 ### Changed diff --git a/README.md b/README.md index 33815ba..0eabcff 100644 --- a/README.md +++ b/README.md @@ -68,11 +68,21 @@ con = ibis.hotdata.connect( default_connection=None, # default catalog (connection id); auto-detected if only one exists default_schema=None, # default schema; auto-detected if only one exists database_id=None, # bind an existing instant database id at connect time + async_after_ms=1000, # queries finishing within this return rows in one request; + # slower ones go async. None = always async poll_interval_s=0.25, # polling interval for async queries poll_timeout_s=600.0, # max time to wait for a query result ) ``` +`api_url`, `token` and `workspace_id` fall back to the `HOTDATA_API_URL`, +`HOTDATA_API_KEY` and `HOTDATA_WORKSPACE` environment variables (same as the +Hotdata CLI), so with those set this is enough: + +```python +con = ibis.hotdata.connect() +``` + URL-style also works, with the same parameters as query string keys: ```python @@ -210,28 +220,48 @@ result = semantic_search(t, "embedding", query_vector, k=10, distance_fn=l2_dist vector column excluded from the output — the SQL shape Hotdata's query engine requires to route the query through its HNSW index instead of a brute-force scan. -Creating the HNSW index itself isn't wrapped by this package yet — use the `hotdata` SDK -directly: +### Text queries + +For a text column indexed with an embedding provider, pass the query as a string. +Hotdata embeds it with the index's provider (`vector_distance(col, 'text')`): ```python -from hotdata import ApiClient, Configuration -from hotdata.api.indexes_api import IndexesApi -from hotdata.models.create_index_request import CreateIndexRequest - -api = IndexesApi(ApiClient(Configuration(...))) -api.create_index( - connection_id=connection_id, - var_schema="main", - table="docs", - create_index_request=CreateIndexRequest( - index_name="docs_embedding_idx", - index_type="vector", - columns=["embedding"], - metric="cosine", - ), +result = semantic_search(t, "body", "how do refunds work?", k=10).execute() +``` + +The text column stays in the output; vector (array) columns such as the generated +`body_embedding` are left out. + +### Managing indexes + +Create, list and drop indexes on instant-database tables (`database=` works as in +`create_table`): + +```python +db = (db_id, "main") + +# HNSW (default): fastest search, memory grows with the table +con.create_index("docs", "embedding", name="docs_idx", database=db, metric="cosine") + +# IVF: for tables too large for HNSW; tune recall/speed with probe_fraction +con.create_index( + "docs", "embedding", name="docs_idx", database=db, + algorithm="ivf", metric="cosine", probe_fraction=0.1, ) + +# Hotdata generates the embeddings from a text column +con.create_index( + "docs", "body", name="docs_body_idx", database=db, + embedding_provider_id="ep_...", +) + +con.list_indexes("docs", database=db) +con.drop_index("docs", "docs_idx", database=db) ``` +`index_type` can also be `"bm25"` (full-text) or `"sorted"`. The index builds within +one request, so raise `timeout=` on `connect` for large tables. + ## Connecting to existing sources If you have existing databases or warehouses connected to your Hotdata workspace (Postgres, Snowflake, BigQuery, etc.), you can query them through the same Ibis connection: @@ -265,13 +295,14 @@ con.list_tables(database=(conn_id, "public")) # tables | `create_database` / `drop_database` (managed) | ✅ | | `create_table` from pandas / PyArrow / schema-only | ✅ | | `drop_table` | ✅ | +| `create_index` / `list_indexes` / `drop_index` (vector HNSW/IVF, BM25, sorted) | ✅ | | `con.table(...)` with full schema metadata | ✅ | | Ibis expressions: filter, select, join, group\_by, agg, order\_by, limit | ✅ | | `con.sql(...)` raw SQL | ✅ | | `.execute()` → pandas, `.to_pyarrow()`, `.to_pyarrow_batches()` | ✅ | | `list_catalogs`, `list_databases`, `list_tables` | ✅ | | Arrow / Parquet column types (timestamp, decimal, list, duration, …) | ✅ | -| Vector search (`ibis_hotdata.vector.semantic_search`) | ✅ | +| Vector search (`ibis_hotdata.vector.semantic_search`, vector or text query) | ✅ | | Temporary tables | ❌ | | In-memory tables (`ibis.memtable(...)`) | ❌ | | Python UDFs | ❌ | diff --git a/examples/06_semantic_search.py b/examples/06_semantic_search.py index 24b07af..52e1181 100644 --- a/examples/06_semantic_search.py +++ b/examples/06_semantic_search.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -"""Vector search: create a managed table with an embedding column, query it with -``ibis_hotdata.vector.semantic_search``. +"""Vector search: create a managed table with an embedding column, index it with +``con.create_index``, query it with ``ibis_hotdata.vector.semantic_search``. Uses small toy 4-dimensional vectors (not real embeddings) split into two obvious clusters -- "pets" and "finance" -- so the nearest-neighbor ordering is easy to @@ -51,10 +51,15 @@ def write(api_url: str, token: str, workspace_id: str) -> str: con = ibis.hotdata.connect(api_url=api_url, token=token, workspace_id=workspace_id) database_id = con.create_database(DATABASE, tables=["docs"], schema=SCHEMA) con.create_table("docs", DOCS, database=(database_id, SCHEMA), overwrite=True) - con.disconnect() - # Uploads are async; give the load a moment to finish before querying. + # Uploads are async; give the load a moment to finish before indexing. time.sleep(2) + con.create_index( + "docs", "embedding", name="docs_embedding_idx", database=(database_id, SCHEMA), + metric="cosine", + ) + print(con.list_indexes("docs", database=(database_id, SCHEMA))) + con.disconnect() return database_id @@ -90,7 +95,7 @@ def main() -> None: token = os.environ["HOTDATA_API_KEY"] workspace_id = os.environ["HOTDATA_WORKSPACE"] - print("== CREATE + WRITE (con.create_database / con.create_table) ==") + print("== CREATE + WRITE + INDEX (create_database / create_table / create_index) ==") database_id = write(API_BASE_URL, token, workspace_id) print(f"created database_id={database_id}") diff --git a/pyproject.toml b/pyproject.toml index a103466..10eae38 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "ibis-framework>=12,<13", - "hotdata>=0.11,<0.12", + "hotdata>=0.12,<0.13", "pyarrow>=16", "pyarrow-hotfix>=0.6", "pandas>=2", diff --git a/src/ibis_hotdata/__init__.py b/src/ibis_hotdata/__init__.py index fd3de97..07b271d 100644 --- a/src/ibis_hotdata/__init__.py +++ b/src/ibis_hotdata/__init__.py @@ -11,6 +11,7 @@ l2_distance, negative_dot_product, semantic_search, + vector_distance, ) __all__ = [ @@ -20,4 +21,5 @@ "l2_distance", "negative_dot_product", "semantic_search", + "vector_distance", ] diff --git a/src/ibis_hotdata/backend.py b/src/ibis_hotdata/backend.py index 83ab5f8..a687ae3 100644 --- a/src/ibis_hotdata/backend.py +++ b/src/ibis_hotdata/backend.py @@ -84,8 +84,11 @@ def _from_url(self, url: ParseResult, **kwarg_overrides: Any): * Base URL defaults to ``https://{host}`` plus optional leading ``path``. * Query string may include ``token``, ``workspace_id``, ``timeout``, ``verify_ssl`` (``true`` / ``false``), ``default_connection``, - ``default_schema``, ``poll_interval_s``, ``poll_timeout_s``. + ``default_schema``, ``database_id``, ``async_after_ms`` (``none`` to + disable), ``poll_interval_s``, ``poll_timeout_s``. * If ``token`` is omitted, ``urlparse`` password (`user:TOKEN@`) is accepted. + * A missing ``token`` / ``workspace_id`` falls back to ``HOTDATA_API_KEY`` / + ``HOTDATA_WORKSPACE``. """ q = dict(parse_qsl(url.query, keep_blank_values=True)) q.update(kwarg_overrides) @@ -111,6 +114,7 @@ def _from_url(self, url: ParseResult, **kwarg_overrides: Any): token = q.pop("token", None) or (unquote_plus(url.password) if url.password else None) workspace_id = q.pop("workspace_id", None) + async_after = q.pop("async_after_ms", "1000") kwargs = { "api_url": api_url, "token": token, @@ -119,45 +123,49 @@ def _from_url(self, url: ParseResult, **kwarg_overrides: Any): "verify_ssl": verify_ssl, "default_connection": q.pop("default_connection", None), "default_schema": q.pop("default_schema", None), + "database_id": q.pop("database_id", None), + "async_after_ms": ( + None if str(async_after).lower() in ("", "none") else int(async_after) + ), "poll_interval_s": float(q.pop("poll_interval_s", "0.25")), "poll_timeout_s": float(q.pop("poll_timeout_s", "600")), } self._convert_kwargs(kwargs) - if not kwargs.get("token"): - raise com.IbisError("Hotdata URL missing token (query parameter or password segment)") - if not kwargs.get("workspace_id"): - raise com.IbisError("Hotdata URL missing workspace_id query parameter") - return self.connect(**kwargs) def do_connect( self, *, - api_url: str, - token: str, - workspace_id: str, + api_url: str | None = None, + token: str | None = None, + workspace_id: str | None = None, timeout: float = 120.0, verify_ssl: bool | str = True, default_connection: str | None = None, default_schema: str | None = None, database_id: str | None = None, + async_after_ms: int | None = 1000, poll_interval_s: float = 0.25, poll_timeout_s: float = 600.0, ) -> None: """Create an Ibis client for a Hotdata workspace. - Query execution always uses Hotdata's async path and downloads ready - results as Arrow IPC from ``GET /v1/results/{id}``. + Queries ask for an inline Arrow IPC response. A query that finishes + within ``async_after_ms`` returns its rows in one request; slower or + larger ones are polled and their results downloaded as Arrow IPC from + ``GET /v1/results/{id}``. Parameters ---------- api_url - Hotdata API base URL (e.g. ``https://api.hotdata.dev``). + Hotdata API base URL. Defaults to ``HOTDATA_API_URL``, else + ``https://api.hotdata.dev``. token - API bearer token (``Authorization`` header). + API bearer token (``Authorization`` header). Defaults to ``HOTDATA_API_KEY``. workspace_id - Workspace public id (``X-Workspace-Id`` header). + Workspace public id (``X-Workspace-Id`` header). Defaults to + ``HOTDATA_WORKSPACE``. timeout HTTP timeout in seconds (per request). verify_ssl @@ -174,6 +182,9 @@ def do_connect( queries use this database's ``X-Database-Id`` header and ``"default"`` catalog references resolve to the underlying connection automatically — without needing to call ``create_table`` first. + async_after_ms + How long (ms, at least 1000) a query may run inline before Hotdata + switches it to async. ``None`` makes every query async. poll_interval_s Sleep between ``GET /v1/query-runs/{id}`` polls. poll_timeout_s @@ -187,14 +198,18 @@ def do_connect( self._database_connection_id: str | None = None self._poll_interval_s = poll_interval_s self._poll_timeout_s = poll_timeout_s + self._async_after_ms = async_after_ms - self._http = HotdataClient( - api_url=api_url, - token=token, - workspace_id=workspace_id, - timeout=timeout, - verify_ssl=verify_ssl, - ) + try: + self._http = HotdataClient( + api_url=api_url, + token=token, + workspace_id=workspace_id, + timeout=timeout, + verify_ssl=verify_ssl, + ) + except HotdataAPIError as exc: + raise _ibis_err_from_hotdata(exc) from exc def disconnect(self) -> None: if hasattr(self, "_http"): @@ -466,6 +481,7 @@ def _get_schema_using_query(self, query: str) -> sch.Schema: data = self._http.execute_query( preview_sql, database_id=self._database_id, + async_after_ms=self._async_after_ms, poll_interval_s=self._poll_interval_s, poll_timeout_s=self._poll_timeout_s, ) @@ -488,6 +504,7 @@ def _safe_raw_sql( payload = self._http.execute_query( query, database_id=self._database_id, + async_after_ms=self._async_after_ms, poll_interval_s=self._poll_interval_s, poll_timeout_s=self._poll_timeout_s, ) @@ -518,6 +535,7 @@ def to_pyarrow( payload = self._http.execute_query( sql, database_id=self._database_id, + async_after_ms=self._async_after_ms, poll_interval_s=self._poll_interval_s, poll_timeout_s=self._poll_timeout_s, ) @@ -726,6 +744,90 @@ def drop_table( return raise _ibis_err_from_hotdata(exc) from exc + def create_index( + self, + table: str, + columns: str | Sequence[str], + /, + *, + name: str, + database: tuple[str, str] | str | None = None, + index_type: str = "vector", + metric: str | None = None, + algorithm: str | None = None, + vector_precision: str | None = None, + nlist: int | None = None, + probe_fraction: float | None = None, + embedding_provider_id: str | None = None, + output_column: str | None = None, + dimensions: int | None = None, + ) -> dict[str, Any]: + """Create an index on an instant-database table and return its details. + + ``index_type`` is ``"vector"`` (default), ``"bm25"`` or ``"sorted"``. + Vector options: ``metric`` (``"l2"``, ``"cosine"``, ``"dot"``), + ``algorithm`` (``"hnsw"`` default, or ``"ivf"`` for tables too large for + HNSW), ``vector_precision``, and for IVF ``nlist`` / ``probe_fraction``. + Set ``embedding_provider_id`` to index a text column: Hotdata generates + the embeddings into ``output_column`` (default ``{column}_embedding``). + + The build runs within one request; raise ``timeout=`` on ``connect`` + for large tables. ``database`` is addressed as in :meth:`create_table`. + """ + cols = [columns] if isinstance(columns, str) else list(columns) + connection_id, schema_name = self._table_location(database) + try: + return self._http.create_index( + connection_id, + schema_name, + table, + index_name=name, + columns=cols, + index_type=index_type, + metric=metric, + algorithm=algorithm, + vector_precision=vector_precision, + nlist=nlist, + probe_fraction=probe_fraction, + embedding_provider_id=embedding_provider_id, + output_column=output_column, + dimensions=dimensions, + ) + except HotdataAPIError as exc: + raise _ibis_err_from_hotdata(exc) from exc + + def list_indexes( + self, + table: str, + /, + *, + database: tuple[str, str] | str | None = None, + ) -> list[dict[str, Any]]: + """List the indexes on an instant-database table.""" + connection_id, schema_name = self._table_location(database) + try: + return self._http.list_indexes(connection_id, schema_name, table) + except HotdataAPIError as exc: + raise _ibis_err_from_hotdata(exc) from exc + + def drop_index( + self, + table: str, + name: str, + /, + *, + database: tuple[str, str] | str | None = None, + force: bool = False, + ) -> None: + """Drop index ``name`` from a table; ``force=True`` ignores a missing index.""" + connection_id, schema_name = self._table_location(database) + try: + self._http.delete_index(connection_id, schema_name, table, name) + except HotdataAPIError as exc: + if force and exc.status_code == 404: + return + raise _ibis_err_from_hotdata(exc) from exc + def _register_in_memory_table(self, _op: ops.InMemoryTable) -> None: pass # Hotdata has no local in-memory table concept; Ibis calls this hook before execute diff --git a/src/ibis_hotdata/http.py b/src/ibis_hotdata/http.py index eedc612..be067da 100644 --- a/src/ibis_hotdata/http.py +++ b/src/ibis_hotdata/http.py @@ -12,19 +12,20 @@ import pyarrow as pa import pyarrow.ipc as pa_ipc import pyarrow_hotfix # noqa: F401 -from hotdata import ApiClient, Configuration + +# Top-level ``ConnectionsApi``/``DatabasesApi`` retry loads while the table is locked. +from hotdata import ApiClient, Configuration, ConnectionsApi, DatabasesApi from hotdata.api import ( - ConnectionsApi, + IndexesApi, InformationSchemaApi, QueryApi, QueryRunsApi, ResultsApi, ) -from hotdata.api.databases_api import DatabasesApi from hotdata.exceptions import ApiException from hotdata.models import QueryRequest -from hotdata.models.async_query_response import AsyncQueryResponse from hotdata.models.create_database_request import CreateDatabaseRequest +from hotdata.models.create_index_request import CreateIndexRequest from hotdata.models.database_default_schema_decl import DatabaseDefaultSchemaDecl from hotdata.models.database_default_table_decl import DatabaseDefaultTableDecl from hotdata.models.load_managed_table_request import LoadManagedTableRequest @@ -71,6 +72,22 @@ def _from_upload_error(exc: UploadError) -> HotdataAPIError: return HotdataAPIError(f"Hotdata upload error: {exc}", status_code=getattr(exc, "status", None)) +def _from_raw_error(status: int, reason: str | None, body: bytes) -> HotdataAPIError: + """Error for a non-success response read without the SDK's deserialization.""" + text = body.decode("utf-8", errors="replace") if body else "" + msg = f"Hotdata API error: {reason or status}" + if text: + msg = f"{msg} {text}" + return HotdataAPIError(msg.strip(), status_code=status, body=text) + + +def _int_header(value: str | None) -> int | None: + try: + return int(value) if value is not None else None + except ValueError: + return None + + def _ipc_stream_bytes_to_table(data: bytes) -> pa.Table: with pa_ipc.open_stream(io.BytesIO(data)) as reader: return reader.read_all() @@ -86,14 +103,22 @@ class HotdataClient: def __init__( self, *, - api_url: str, - token: str, - workspace_id: str, + api_url: str | None = None, + token: str | None = None, + workspace_id: str | None = None, timeout: float = 120.0, verify_ssl: bool | str = True, ) -> None: - host = api_url.rstrip("/") + # Unset values fall back to HOTDATA_API_URL / HOTDATA_API_KEY / + # HOTDATA_WORKSPACE (resolved by the SDK's Configuration). + host = api_url.rstrip("/") if api_url else None conf = Configuration(host=host, api_key=token, workspace_id=workspace_id) + if not conf.api_key: + raise HotdataAPIError("Missing Hotdata token: pass token= or set HOTDATA_API_KEY") + if not conf.api_keys.get("WorkspaceId"): + raise HotdataAPIError( + "Missing Hotdata workspace: pass workspace_id= or set HOTDATA_WORKSPACE" + ) if verify_ssl is False: conf.verify_ssl = False elif isinstance(verify_ssl, str): @@ -106,6 +131,7 @@ def __init__( self._connections = ConnectionsApi(self._client) self._databases = DatabasesApi(self._client) self._information_schema = InformationSchemaApi(self._client) + self._indexes = IndexesApi(self._client) self._uploads = UploadsApi(self._client) def close(self) -> None: @@ -150,43 +176,104 @@ def execute_query( sql: str, *, database_id: str | None = None, - async_after_ms: int | None = None, + async_after_ms: int | None = 1000, poll_interval_s: float = 0.25, poll_timeout_s: float = 600.0, ) -> dict[str, Any]: - req = QueryRequest(sql=sql, var_async=True, async_after_ms=async_after_ms) + """Run ``sql`` and return its result as an Arrow payload. + + Asks for an inline Arrow IPC response: a query that finishes within + ``async_after_ms`` returns its rows in this one request. A slower or + larger query falls back to polling the query run and its result. + ``async_after_ms=None`` skips the inline attempt (always async). + """ + req = QueryRequest(sql=sql, var_async=True) + if async_after_ms is not None: + req.async_after_ms = async_after_ms kwargs: dict[str, Any] = {} if database_id is not None: kwargs["x_database_id"] = database_id - out = self._safe_call(self._query.query, req, **kwargs) - if isinstance(out, AsyncQueryResponse): - query_run_id = out.query_run_id - effective_database_id = database_id or "" - deadline = time.monotonic() + poll_timeout_s - while time.monotonic() < deadline: - qr = self._safe_call( - self._query_runs.get_query_run, + try: + raw = self._query.query_without_preload_content( + req, + _headers={"Accept": APPLICATION_ARROW_STREAM}, + _request_timeout=self._timeout, + **kwargs, + ) + except ApiException as exc: + raise _from_api_exception(exc) from exc + body = raw.read() + status = raw.status + headers = raw.headers + ctype = (headers.get("Content-Type") or "").split(";")[0].strip().lower() + effective_database_id = database_id or "" + deadline = time.monotonic() + poll_timeout_s + + if status == http.HTTPStatus.OK and ctype == APPLICATION_ARROW_STREAM.lower(): + return self._arrow_payload_from_table( + _ipc_stream_bytes_to_table(body), + result_id=headers.get("X-Result-Id"), + query_run_id=headers.get("X-Query-Run-Id"), + execution_time_ms=_int_header(headers.get("X-Execution-Time-Ms")), + warning=headers.get("X-Query-Warning"), + ) + + if status in (http.HTTPStatus.OK, http.HTTPStatus.ACCEPTED): + d = _json_utf8(body) if body else {} + # Result too large to inline: the server names the result to fetch. + result_id = headers.get("X-Result-Id") or ( + d.get("result_id") if status == http.HTTPStatus.OK else None + ) + if result_id: + return self._poll_result_arrow( + result_id, + database_id=effective_database_id, + deadline=deadline, + poll_interval_s=poll_interval_s, + ) + query_run_id = d.get("query_run_id") + if query_run_id: + return self._poll_query_run( query_run_id, - x_database_id=effective_database_id, + database_id=effective_database_id, + deadline=deadline, + poll_interval_s=poll_interval_s, ) - status = qr.status - if status == "failed": - raise HotdataAPIError(qr.error_message or "Query run failed") - if status == "succeeded": - result_id = qr.result_id - if result_id is None: - raise HotdataAPIError("succeeded query run missing result_id") - return self._poll_result_arrow( - result_id, - database_id=effective_database_id, - deadline=deadline, - poll_interval_s=poll_interval_s, - ) - if status not in _IN_FLIGHT: - raise HotdataAPIError(f"Unexpected query run status: {status!r}") - _sleep_until(deadline, poll_interval_s) - raise HotdataAPIError("Timeout waiting for asynchronous query") - raise HotdataAPIError("Unexpected query response type") + raise HotdataAPIError("Unexpected query response type", status_code=status, body=d) + + raise _from_raw_error(status, raw.reason, body) + + def _poll_query_run( + self, + query_run_id: str, + *, + database_id: str, + deadline: float, + poll_interval_s: float, + ) -> dict[str, Any]: + while time.monotonic() < deadline: + qr = self._safe_call( + self._query_runs.get_query_run, + query_run_id, + x_database_id=database_id, + ) + status = qr.status + if status == "failed": + raise HotdataAPIError(qr.error_message or "Query run failed") + if status == "succeeded": + result_id = qr.result_id + if result_id is None: + raise HotdataAPIError("succeeded query run missing result_id") + return self._poll_result_arrow( + result_id, + database_id=database_id, + deadline=deadline, + poll_interval_s=poll_interval_s, + ) + if status not in _IN_FLIGHT: + raise HotdataAPIError(f"Unexpected query run status: {status!r}") + _sleep_until(deadline, poll_interval_s) + raise HotdataAPIError("Timeout waiting for asynchronous query") def upload_file(self, data: bytes, *, content_type: str | None = None) -> dict[str, Any]: """Direct-to-storage presigned upload: create session, ``PUT``, finalize. @@ -267,6 +354,37 @@ def delete_managed_table(self, connection_id: str, schema: str, table: str) -> N table, ) + def create_index( + self, + connection_id: str, + schema: str, + table: str, + *, + index_name: str, + columns: Sequence[str], + **options: Any, + ) -> dict[str, Any]: + """POST ``.../tables/{table}/indexes`` (synchronous build). + + ``options`` are ``CreateIndexRequest`` fields; ``None`` values are dropped. + """ + fields = {k: v for k, v in options.items() if v is not None} + try: + req = CreateIndexRequest(index_name=index_name, columns=list(columns), **fields) + except ValueError as exc: # pydantic: unknown field or invalid value + raise HotdataAPIError(f"Invalid index request: {exc}") from exc + resp = self._safe_call(self._indexes.create_index, connection_id, schema, table, req) + return resp.model_dump(by_alias=True, mode="json") + + def list_indexes(self, connection_id: str, schema: str, table: str) -> list[dict[str, Any]]: + """GET ``.../tables/{table}/indexes``.""" + resp = self._safe_call(self._indexes.list_indexes, connection_id, schema, table) + return [i.model_dump(by_alias=True, mode="json") for i in resp.indexes] + + def delete_index(self, connection_id: str, schema: str, table: str, index_name: str) -> None: + """DELETE ``.../tables/{table}/indexes/{index_name}``.""" + self._safe_call(self._indexes.delete_index, connection_id, schema, table, index_name) + def _poll_result_arrow( self, result_id: str, @@ -326,7 +444,10 @@ def _arrow_payload_from_table( self, table: pa.Table, *, - result_id: str, + result_id: str | None, + query_run_id: str | None = None, + execution_time_ms: int | None = None, + warning: str | None = None, ) -> dict[str, Any]: sch = table.schema columns = sch.names @@ -339,7 +460,7 @@ def _arrow_payload_from_table( "rows": [], "result_id": result_id, "row_count": table.num_rows, - "execution_time_ms": None, - "query_run_id": None, - "warning": None, + "execution_time_ms": execution_time_ms, + "query_run_id": query_run_id, + "warning": warning, } diff --git a/src/ibis_hotdata/vector.py b/src/ibis_hotdata/vector.py index 1f8dbff..eb3c39a 100644 --- a/src/ibis_hotdata/vector.py +++ b/src/ibis_hotdata/vector.py @@ -1,6 +1,6 @@ """Vector-search helpers: distance UDFs + a semantic-search query builder. -Bridges Hotdata's HNSW-indexed distance functions into Ibis. These are pure +Bridges Hotdata's vector-indexed (HNSW / IVF) distance functions into Ibis. These are pure functions over Ibis expressions built on ``@ibis.udf.scalar.builtin`` — no backend changes required, and no Python is shipped to the engine. """ @@ -28,10 +28,18 @@ def negative_dot_product(a, b) -> float: """Negative dot product between two vectors (smaller is more similar).""" +@ibis.udf.scalar.builtin +def vector_distance(a, b) -> float: + """Distance between a text column with an embedding-backed index and a text query. + + Hotdata embeds the query with the index's provider and uses the index's metric. + """ + + def semantic_search( table: ir.Table, - column: str | ir.ArrayColumn, - query_vector: Sequence[float], + column: str | ir.Column, + query_vector: Sequence[float] | str, k: int, *, distance_fn=cosine_distance, @@ -39,17 +47,27 @@ def semantic_search( ) -> ir.Table: """Return the `k` rows of `table` whose `column` is nearest `query_vector`. - Excludes `column` from the result and orders ascending by - `distance_fn(column, query_vector)`, aliased as `distance_name`. + With a vector query, excludes `column` from the result and orders + ascending by `distance_fn(column, query_vector)`, aliased as `distance_name`. + + With a text query (a `str`), `column` is the source text column of an + embedding-backed index: orders by `vector_distance(column, query)` + (`distance_fn` is ignored) and excludes every array (vector) column, so + the generated embedding column is left out of the result. `distance_name` must not collide with an existing column name in `table`, or `select` raises a duplicate-column error. """ col = table[column] if isinstance(column, str) else column - other_cols = [name for name in table.columns if name != col.get_name()] - qvec = ibis.literal(list(query_vector)) + if isinstance(query_vector, str): + schema = table.schema() + other_cols = [name for name in table.columns if not schema[name].is_array()] + dist = vector_distance(col, ibis.literal(query_vector)) + else: + other_cols = [name for name in table.columns if name != col.get_name()] + dist = distance_fn(col, ibis.literal(list(query_vector))) return ( - table.select(*other_cols, **{distance_name: distance_fn(col, qvec)}) + table.select(*other_cols, **{distance_name: dist}) .order_by(ibis.asc(distance_name)) .limit(k) ) diff --git a/tests/test_hotdata_backend.py b/tests/test_hotdata_backend.py index 230ee2b..bd364c7 100644 --- a/tests/test_hotdata_backend.py +++ b/tests/test_hotdata_backend.py @@ -37,12 +37,13 @@ ] -# TableInfo requires these from hotdata 0.9.0 on. The API always sends them, and -# an omission fails validation for the WHOLE information_schema call rather than -# for that field — nine tests here failed exactly that way before it was added. +# TableInfo requires these from hotdata 0.9.0 on (constant_per_key from 0.12.0). +# The API always sends them, and an omission fails validation for the WHOLE +# information_schema call rather than for that field — nine tests here failed +# exactly that way before each was added. # Spread into each fixture rather than repeated, so the next required field is a # one-line change here. -_REQUIRED_TABLE_FIELDS = {"partition_by": [], "sorted_by": []} +_REQUIRED_TABLE_FIELDS = {"constant_per_key": [], "partition_by": [], "sorted_by": []} def arrow_stream(table: pa.Table) -> bytes: @@ -876,3 +877,97 @@ def on_post(req: Request) -> Response: ibis.hotdata.connect( api_url=srv, token="tok", workspace_id="ws", session_id="sb_xyz" ) + + +def test_connect_reads_credentials_from_env(httpserver: HTTPServer, srv: str, monkeypatch): + monkeypatch.setenv("HOTDATA_API_URL", srv) + monkeypatch.setenv("HOTDATA_API_KEY", "envtok") + monkeypatch.setenv("HOTDATA_WORKSPACE", "ws_env") + httpserver.expect_request("/v1/connections").respond_with_json( + {"connections": [{"id": TPCH_CONN, "name": "tpch", "source_type": "duckdb"}]} + ) + + con = ibis.hotdata.connect(verify_ssl=False) + assert con.list_catalogs() == [TPCH_CONN] + + +def test_connect_without_credentials_raises(monkeypatch): + monkeypatch.delenv("HOTDATA_API_KEY", raising=False) + monkeypatch.setenv("HOTDATA_WORKSPACE", "ws_env") + with pytest.raises(com.IbisError, match="HOTDATA_API_KEY"): + ibis.hotdata.connect() + + +@pytest.mark.parametrize(("value", "expected"), [("5000", 5000), ("none", None)]) +def test_url_async_after_ms(httpserver: HTTPServer, srv: str, value, expected): + url = ( + f"hotdata://127.0.0.1:{httpserver.port}?token=tok&workspace_id=ws" + f"&verify_ssl=false&api_url={srv}&async_after_ms={value}" + ) + con = ibis.connect(url) + assert con._async_after_ms == expected + + +_INDEX_PATH = f"/v1/connections/{MANAGED_CONN}/tables/{PUBLIC}/docs/indexes" +_INDEX_INFO = { + "index_name": "docs_idx", + "index_type": "vector", + "columns": ["embedding"], + "algorithm": "hnsw", + "metric": "cosine", + "status": "ready", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", +} + + +def test_create_index_resolves_instant_database(httpserver: HTTPServer, srv: str): + httpserver.expect_request(f"/v1/databases/{MANAGED_DB_ID}").respond_with_json( + managed_database_detail_response() + ) + + def on_create(req: Request) -> Response: + assert req.get_json() == { + "index_name": "docs_idx", + "columns": ["embedding"], + "index_type": "vector", + "metric": "cosine", + "async": False, + } + return Response(json.dumps(_INDEX_INFO), status=201, content_type="application/json") + + httpserver.expect_oneshot_request(_INDEX_PATH, method="POST").respond_with_handler(on_create) + + con = ibis.hotdata.connect(api_url=srv, token="tok", workspace_id="ws", verify_ssl=False) + out = con.create_index( + "docs", "embedding", name="docs_idx", database=(MANAGED_DB_ID, PUBLIC), metric="cosine" + ) + assert out["index_name"] == "docs_idx" + httpserver.check_assertions() + + +def test_list_indexes(httpserver: HTTPServer, srv: str): + httpserver.expect_request(f"/v1/databases/{MANAGED_DB_ID}").respond_with_json( + managed_database_detail_response() + ) + httpserver.expect_request(_INDEX_PATH, method="GET").respond_with_json( + {"indexes": [_INDEX_INFO]} + ) + + con = ibis.hotdata.connect(api_url=srv, token="tok", workspace_id="ws", verify_ssl=False) + out = con.list_indexes("docs", database=(MANAGED_DB_ID, PUBLIC)) + assert [i["index_name"] for i in out] == ["docs_idx"] + + +def test_drop_index_force_ignores_missing(httpserver: HTTPServer, srv: str): + httpserver.expect_request(f"/v1/databases/{MANAGED_DB_ID}").respond_with_json( + managed_database_detail_response() + ) + httpserver.expect_request(f"{_INDEX_PATH}/nope", method="DELETE").respond_with_data( + b"not found", status=404 + ) + + con = ibis.hotdata.connect(api_url=srv, token="tok", workspace_id="ws", verify_ssl=False) + con.drop_index("docs", "nope", database=(MANAGED_DB_ID, PUBLIC), force=True) + with pytest.raises(com.IbisError): + con.drop_index("docs", "nope", database=(MANAGED_DB_ID, PUBLIC)) diff --git a/tests/test_hotdata_http.py b/tests/test_hotdata_http.py index 18ee717..4e30e63 100644 --- a/tests/test_hotdata_http.py +++ b/tests/test_hotdata_http.py @@ -354,3 +354,176 @@ def on_create_session(req: Request) -> Response: out = client.upload_file(b"parquet", content_type="application/parquet") assert out["upload_id"] == "upl_1" client.close() + + +def _arrow_blob(table: pa.Table) -> bytes: + sink = io.BytesIO() + with ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return sink.getvalue() + + +def _client(httpserver: HTTPServer) -> HotdataClient: + return HotdataClient( + api_url=httpserver.url_for("/").rstrip("/"), + token="tok", + workspace_id="ws1", + verify_ssl=False, + ) + + +def test_execute_query_inline_arrow_single_request(httpserver: HTTPServer): + """A query that finishes in time returns its rows in the POST response itself.""" + + def on_query(req: Request) -> Response: + assert req.headers["Accept"] == APPLICATION_ARROW_STREAM + assert req.get_json() == {"async": True, "async_after_ms": 1000, "sql": "select 42"} + return Response( + _arrow_blob(pa.table({"n": [42]})), + status=200, + content_type=APPLICATION_ARROW_STREAM, + headers={ + "X-Query-Run-Id": "run1", + "X-Result-Id": "res1", + "X-Execution-Time-Ms": "7", + "X-Query-Warning": "careful", + }, + ) + + httpserver.expect_oneshot_request("/v1/query", method="POST").respond_with_handler(on_query) + + client = _client(httpserver) + body = client.execute_query("select 42") + client.close() + + assert body["pa_table"].to_pydict() == {"n": [42]} + assert body["result_id"] == "res1" + assert body["query_run_id"] == "run1" + assert body["execution_time_ms"] == 7 + assert body["warning"] == "careful" + httpserver.check_assertions() + + +def test_execute_query_async_after_ms_none_sends_pure_async(httpserver: HTTPServer): + def on_query(req: Request) -> Response: + assert req.get_json() == {"async": True, "sql": "select 1"} + return Response( + _arrow_blob(pa.table({"n": [1]})), status=200, content_type=APPLICATION_ARROW_STREAM + ) + + httpserver.expect_oneshot_request("/v1/query", method="POST").respond_with_handler(on_query) + client = _client(httpserver) + assert client.execute_query("select 1", async_after_ms=None)["pa_table"].num_rows == 1 + client.close() + + +def test_execute_query_too_large_to_inline_fetches_named_result(httpserver: HTTPServer): + """202 with X-Result-Id: read the result directly, skipping the query-run poll.""" + httpserver.expect_oneshot_request("/v1/query", method="POST").respond_with_json( + {"query_run_id": "run1", "status": "running", "status_url": "http://poll"}, + status=202, + headers={"X-Result-Id": "res1"}, + ) + httpserver.expect_oneshot_request("/v1/results/res1").respond_with_data( + _arrow_blob(pa.table({"n": [1, 2]})), status=200, content_type=APPLICATION_ARROW_STREAM + ) + + client = _client(httpserver) + body = client.execute_query("select n", poll_interval_s=0, poll_timeout_s=5) + client.close() + + assert body["pa_table"].to_pydict() == {"n": [1, 2]} + httpserver.check_assertions() + + +def test_execute_query_inline_error_carries_status(httpserver: HTTPServer): + httpserver.expect_oneshot_request("/v1/query", method="POST").respond_with_json( + {"error": {"message": "syntax error"}}, status=400 + ) + client = _client(httpserver) + with pytest.raises(HotdataAPIError, match="syntax error") as info: + client.execute_query("selec 1") + client.close() + assert info.value.status_code == 400 + + +def test_client_reads_credentials_from_env(httpserver: HTTPServer, monkeypatch): + monkeypatch.setenv("HOTDATA_API_URL", httpserver.url_for("/").rstrip("/")) + monkeypatch.setenv("HOTDATA_API_KEY", "envtok") + monkeypatch.setenv("HOTDATA_WORKSPACE", "ws_env") + + def on_query(req: Request) -> Response: + assert req.headers["X-Workspace-Id"] == "ws_env" + return Response( + _arrow_blob(pa.table({"n": [1]})), status=200, content_type=APPLICATION_ARROW_STREAM + ) + + httpserver.expect_oneshot_request("/v1/query", method="POST").respond_with_handler(on_query) + client = HotdataClient(verify_ssl=False) + assert client.execute_query("select 1")["pa_table"].num_rows == 1 + client.close() + httpserver.check_assertions() + + +@pytest.mark.parametrize( + ("missing", "match"), + [("HOTDATA_API_KEY", "HOTDATA_API_KEY"), ("HOTDATA_WORKSPACE", "HOTDATA_WORKSPACE")], +) +def test_client_missing_credentials_raises(monkeypatch, missing, match): + monkeypatch.setenv("HOTDATA_API_KEY", "envtok") + monkeypatch.setenv("HOTDATA_WORKSPACE", "ws_env") + monkeypatch.delenv(missing) + with pytest.raises(HotdataAPIError, match=match): + HotdataClient() + + +_INDEX = { + "index_name": "docs_idx", + "index_type": "vector", + "columns": ["embedding"], + "algorithm": "ivf", + "metric": "cosine", + "status": "ready", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", +} + + +def test_create_index_drops_unset_options(httpserver: HTTPServer): + def on_create(req: Request) -> Response: + assert req.get_json() == { + "index_name": "docs_idx", + "columns": ["embedding"], + "index_type": "vector", + "algorithm": "ivf", + "metric": "cosine", + "async": False, + } + return Response(json.dumps(_INDEX), status=201, content_type="application/json") + + httpserver.expect_oneshot_request( + "/v1/connections/conn1/tables/main/docs/indexes", method="POST" + ).respond_with_handler(on_create) + + client = _client(httpserver) + out = client.create_index( + "conn1", + "main", + "docs", + index_name="docs_idx", + columns=["embedding"], + index_type="vector", + algorithm="ivf", + metric="cosine", + nlist=None, + ) + client.close() + assert out["algorithm"] == "ivf" + httpserver.check_assertions() + + +def test_create_index_rejects_invalid_option_before_sending(httpserver: HTTPServer): + client = _client(httpserver) + with pytest.raises(HotdataAPIError, match="Invalid index request"): + client.create_index("c", "s", "t", index_name="i", columns=["x"], algorithm="bogus") + client.close() diff --git a/tests/test_vector.py b/tests/test_vector.py index 02a9694..c59b2e5 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -8,6 +8,7 @@ l2_distance, negative_dot_product, semantic_search, + vector_distance, ) ITEMS = ibis.table( @@ -73,3 +74,25 @@ def test_semantic_search_honors_custom_distance_fn_and_name(): def test_semantic_search_accepts_column_expression(): expr = semantic_search(ITEMS, ITEMS.embedding, [0.1, 0.2, 0.3], k=5) assert "embedding" not in expr.columns + + +DOCS = ibis.table( + {"id": "int64", "body": "string", "body_embedding": "array"}, + name="docs", +) + + +def test_vector_distance_compiles_to_function_call(): + expr = DOCS.select(dist=vector_distance(DOCS.body, ibis.literal("cats"))) + sql = ibis.to_sql(expr, dialect="postgres") + assert "VECTOR_DISTANCE(\"t0\".\"body\", 'cats')" in sql + + +def test_semantic_search_text_query_uses_vector_distance(): + expr = semantic_search(DOCS, "body", "cats", k=3, distance_fn=l2_distance) + sql = ibis.to_sql(expr, dialect="postgres") + assert "VECTOR_DISTANCE" in sql + assert "L2_DISTANCE" not in sql + assert "LIMIT 3" in sql + # Keeps the text column, drops the generated vector column. + assert set(expr.columns) == {"id", "body", "distance"} diff --git a/uv.lock b/uv.lock index 9548963..5f784f0 100644 --- a/uv.lock +++ b/uv.lock @@ -127,7 +127,7 @@ wheels = [ [[package]] name = "hotdata" -version = "0.11.0" +version = "0.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pydantic" }, @@ -135,9 +135,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1d/10/1a21c58de1fee8c63bf983db9531704e9b0611c63fd7a849a9ffdf9cc8c7/hotdata-0.11.0.tar.gz", hash = "sha256:a0e9fd242baa51b9350a969222319bcdf2c6ede1cf2e8a99e3f81d608af7b6fc", size = 207167, upload-time = "2026-09-18T20:58:41.802Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f2/f6/b13b210543e5af7984ca144e114f35ff02336a7b183ed1436c3846878337/hotdata-0.12.0.tar.gz", hash = "sha256:a78c426e2005ef92bc10152aa808a7d07c760f30508f5d1be77e97fccb2faa11", size = 240493, upload-time = "2026-10-15T19:43:47.934Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/a4/cc71105c5971d29b40cfff76ab592fed95e1d433903631f9d2fccede50a4/hotdata-0.11.0-py3-none-any.whl", hash = "sha256:4ece2f43bac9a4f43ad32c305aa93042aa4a4f68a5bbdec4f2c6ae0b29f591d6", size = 292742, upload-time = "2026-09-18T20:58:40.049Z" }, + { url = "https://files.pythonhosted.org/packages/9a/b8/c00500536b592e496a9c9d68551ead8177e0dda06bce48bd137f5c861f7f/hotdata-0.12.0-py3-none-any.whl", hash = "sha256:812d96ad6730eac1ca87c6844b89318104fb73068b99f4fae9375f5ea5835217", size = 330254, upload-time = "2026-10-15T19:43:45.995Z" }, ] [[package]] @@ -165,7 +165,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "hotdata", specifier = ">=0.11,<0.12" }, + { name = "hotdata", specifier = ">=0.12,<0.13" }, { name = "ibis-framework", specifier = ">=12,<13" }, { name = "pandas", specifier = ">=2" }, { name = "pyarrow", specifier = ">=16" }, @@ -1214,12 +1214,12 @@ wheels = [ [[package]] name = "werkzeug" -version = "3.1.8" +version = "3.1.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markupsafe" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/dd/b2/381be8cfdee792dd117872481b6e378f85c957dd7c5bca38897b08f765fd/werkzeug-3.1.8.tar.gz", hash = "sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44", size = 875852, upload-time = "2026-04-02T18:49:14.268Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a4/34/4dd12fc8bb7d61c91467ec3efe415ffa7d5456f799954b40c5bbaeae470e/werkzeug-3.1.9.tar.gz", hash = "sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060", size = 940188, upload-time = "2026-09-27T18:33:41.637Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50", size = 226459, upload-time = "2026-04-02T18:49:12.72Z" }, + { url = "https://files.pythonhosted.org/packages/a1/38/df03f564f43cec2684823f3cccae1a652ee7face1cbaa76fb223096e64d7/werkzeug-3.1.9-py3-none-any.whl", hash = "sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab", size = 228700, upload-time = "2026-09-27T18:33:39.685Z" }, ]