From 981dc4d568e06469a52d781e35090968c3fce3e3 Mon Sep 17 00:00:00 2001 From: Rohan Dsouza Date: Tue, 21 Jul 2026 18:00:55 +0530 Subject: [PATCH 1/2] feat: add vector-search helpers (cosine/l2/dot distance UDFs, semantic_search) Adds ibis_hotdata.vector: three @ibis.udf.scalar.builtin distance-function stubs plus a semantic_search(table, column, query_vector, k) query builder, compiling to ORDER BY (...) ASC LIMIT k with the vector column excluded from output -- the SQL shape the engine's HNSW index-selection rule requires. Verified via ibis.to_sql compile-only tests, and live against a local-cluster workspace (examples/06_semantic_search.py). --- CHANGELOG.md | 7 +++ README.md | 43 ++++++++++++++ examples/06_semantic_search.py | 105 +++++++++++++++++++++++++++++++++ src/ibis_hotdata/__init__.py | 15 ++++- src/ibis_hotdata/vector.py | 52 ++++++++++++++++ tests/test_vector.py | 75 +++++++++++++++++++++++ 6 files changed, 296 insertions(+), 1 deletion(-) create mode 100644 examples/06_semantic_search.py create mode 100644 src/ibis_hotdata/vector.py create mode 100644 tests/test_vector.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 59dd82e..5bc0f58 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `ibis_hotdata.vector`: `cosine_distance`, `l2_distance`, `negative_dot_product` + builtin-UDF helpers and a `semantic_search(table, column, query_vector, k, ...)` + query builder, for querying HNSW-indexed vector columns. Compiles to + `ORDER BY (col, ARRAY[...]) ASC LIMIT k` with the vector column excluded + from the output, which is the SQL shape the engine's index-selection rule requires. ## [0.3.2] - 2026-07-20 diff --git a/README.md b/README.md index 5b47f06..11c4376 100644 --- a/README.md +++ b/README.md @@ -191,6 +191,48 @@ result = base.filter(base.amount > 10).execute() You can chain Ibis expressions on the result of `con.sql(...)`. +## Vector search + +`ibis_hotdata.vector` provides helpers for querying HNSW-indexed vector (embedding) +columns: + +```python +from ibis_hotdata.vector import semantic_search, l2_distance + +t = con.table("docs", database=("default", "main")) + +result = semantic_search(t, "embedding", query_vector, k=10).execute() + +# or with a different metric +result = semantic_search(t, "embedding", query_vector, k=10, distance_fn=l2_distance).execute() +``` + +`semantic_search` compiles to `ORDER BY (col, ARRAY[...]) ASC LIMIT k` with the +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: + +```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", + ), +) +``` + ## 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: @@ -228,6 +270,7 @@ con.list_tables(database=("my_postgres", "public")) # tables | `.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`) | ✅ | | Temporary tables | ❌ | | In-memory tables (`ibis.memtable(...)`) | ❌ | | Python UDFs | ❌ | diff --git a/examples/06_semantic_search.py b/examples/06_semantic_search.py new file mode 100644 index 0000000..24b07af --- /dev/null +++ b/examples/06_semantic_search.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Vector search: create a managed table with an embedding column, 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 +eyeball. Demonstrates that the compiled SQL keeps the embedding column out of the +result set and orders by the aliased distance column ascending, which is the shape +Hotdata's query engine requires to route the query through its HNSW index rather +than falling back to a brute-force scan. + +Run against hosted Hotdata (the default) or a local cluster: + + HOTDATA_API_KEY=... HOTDATA_WORKSPACE=... \\ + uv run python examples/06_semantic_search.py + +Point at a local cluster instead by exporting HOTDATA_API_BASE_URL=http://api.localhost +(plus a local key/workspace) before running. + +Env: + HOTDATA_API_KEY, HOTDATA_WORKSPACE -- required + HOTDATA_API_BASE_URL -- optional (default https://api.hotdata.dev) +""" + +from __future__ import annotations + +import os +import time + +import ibis +import pandas as pd + +from ibis_hotdata.vector import l2_distance, semantic_search + +DATABASE = "ibis_semantic_search_demo" +SCHEMA = "public" +API_BASE_URL = os.environ.get("HOTDATA_API_BASE_URL", "https://api.hotdata.dev") + +DOCS = pd.DataFrame( + [ + {"doc_id": "d1", "text": "cats are great pets", "embedding": [0.90, 0.10, 0.00, 0.00]}, + {"doc_id": "d2", "text": "dogs are loyal companions", "embedding": [0.85, 0.15, 0.05, 0.00]}, + {"doc_id": "d3", "text": "stock market rose today", "embedding": [0.00, 0.00, 0.90, 0.10]}, + {"doc_id": "d4", "text": "interest rates and inflation", "embedding": [0.05, 0.00, 0.85, 0.10]}, + ] +) +QUERY_VECTOR = [0.88, 0.12, 0.00, 0.00] # closest to the "pets" cluster (d1, d2) + + +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. + time.sleep(2) + return database_id + + +def query(api_url: str, token: str, workspace_id: str, database_id: str) -> None: + con = ibis.hotdata.connect( + api_url=api_url, + token=token, + workspace_id=workspace_id, + default_schema=SCHEMA, + database_id=database_id, + ) + t = con.table("docs", database=("default", SCHEMA)) + + print("-- semantic_search (cosine, default), compiled SQL --") + cosine_expr = semantic_search(t, "embedding", QUERY_VECTOR, k=3) + print(con.compile(cosine_expr)) + print(cosine_expr.execute()) + + print("\n-- semantic_search with l2_distance instead --") + l2_expr = semantic_search(t, "embedding", QUERY_VECTOR, k=3, distance_fn=l2_distance) + print(l2_expr.execute()) + + con.disconnect() + + +def cleanup(api_url: str, token: str, workspace_id: str, database_id: str) -> None: + con = ibis.hotdata.connect(api_url=api_url, token=token, workspace_id=workspace_id) + con.drop_database(database_id, force=True) + con.disconnect() + + +def main() -> None: + token = os.environ["HOTDATA_API_KEY"] + workspace_id = os.environ["HOTDATA_WORKSPACE"] + + print("== CREATE + WRITE (con.create_database / con.create_table) ==") + database_id = write(API_BASE_URL, token, workspace_id) + print(f"created database_id={database_id}") + + print("\n== QUERY via semantic_search ==") + query(API_BASE_URL, token, workspace_id, database_id) + + print("\n== CLEANUP ==") + cleanup(API_BASE_URL, token, workspace_id, database_id) + + +if __name__ == "__main__": + main() diff --git a/src/ibis_hotdata/__init__.py b/src/ibis_hotdata/__init__.py index 8b5f13b..fd3de97 100644 --- a/src/ibis_hotdata/__init__.py +++ b/src/ibis_hotdata/__init__.py @@ -6,5 +6,18 @@ __version__ = "0.0.0+unknown" from ibis_hotdata.backend import Backend +from ibis_hotdata.vector import ( + cosine_distance, + l2_distance, + negative_dot_product, + semantic_search, +) -__all__ = ["Backend", "__version__"] +__all__ = [ + "Backend", + "__version__", + "cosine_distance", + "l2_distance", + "negative_dot_product", + "semantic_search", +] diff --git a/src/ibis_hotdata/vector.py b/src/ibis_hotdata/vector.py new file mode 100644 index 0000000..de0d142 --- /dev/null +++ b/src/ibis_hotdata/vector.py @@ -0,0 +1,52 @@ +"""Vector-search helpers: distance UDFs + a semantic-search query builder. + +Bridges Hotdata's HNSW-indexed 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. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import ibis +import ibis.expr.types as ir + + +@ibis.udf.scalar.builtin +def cosine_distance(a, b) -> float: + """Cosine distance between two vectors.""" + + +@ibis.udf.scalar.builtin +def l2_distance(a, b) -> float: + """Euclidean (L2) distance between two vectors.""" + + +@ibis.udf.scalar.builtin +def negative_dot_product(a, b) -> float: + """Negative dot product between two vectors (smaller is more similar).""" + + +def semantic_search( + table: ir.Table, + column: str | ir.ArrayColumn, + query_vector: Sequence[float], + k: int, + *, + distance_fn=cosine_distance, + distance_name: str = "distance", +) -> 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`. + """ + 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)) + return ( + table.select(*other_cols, **{distance_name: distance_fn(col, qvec)}) + .order_by(ibis.asc(distance_name)) + .limit(k) + ) diff --git a/tests/test_vector.py b/tests/test_vector.py new file mode 100644 index 0000000..02a9694 --- /dev/null +++ b/tests/test_vector.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import ibis +import pytest + +from ibis_hotdata.vector import ( + cosine_distance, + l2_distance, + negative_dot_product, + semantic_search, +) + +ITEMS = ibis.table( + {"id": "int64", "name": "string", "embedding": "array"}, + name="items", +) + + +@pytest.mark.parametrize( + ("fn", "expected_call"), + [ + (cosine_distance, "COSINE_DISTANCE"), + (l2_distance, "L2_DISTANCE"), + (negative_dot_product, "NEGATIVE_DOT_PRODUCT"), + ], +) +def test_distance_udf_compiles_to_function_call(fn, expected_call): + expr = ITEMS.select(dist=fn(ITEMS.embedding, ibis.literal([0.1, 0.2, 0.3]))) + sql = ibis.to_sql(expr, dialect="postgres") + assert f'{expected_call}("t0"."embedding", ARRAY[0.1, 0.2, 0.3])' in sql + + +def test_semantic_search_query_vector_is_a_literal_array(): + expr = semantic_search(ITEMS, "embedding", [0.1, 0.2, 0.3], k=5) + sql = ibis.to_sql(expr, dialect="postgres") + assert "ARRAY[0.1, 0.2, 0.3]" in sql + + +def test_semantic_search_aliases_distance_and_excludes_source_column(): + expr = semantic_search(ITEMS, "embedding", [0.1, 0.2, 0.3], k=5) + sql = ibis.to_sql(expr, dialect="postgres") + assert 'AS "distance"' in sql + # The embedding column is a valid distance-function argument, but must not + # appear as an output column (engine issue #508: vector-in-output disables + # the HNSW fast path). + assert set(expr.columns) == {"id", "name", "distance"} + + +def test_semantic_search_orders_ascending_with_limit(): + expr = semantic_search(ITEMS, "embedding", [0.1, 0.2, 0.3], k=5) + sql = ibis.to_sql(expr, dialect="postgres") + assert "ORDER BY" in sql + assert '"distance" ASC' in sql + assert "LIMIT 5" in sql + + +def test_semantic_search_honors_custom_distance_fn_and_name(): + expr = semantic_search( + ITEMS, + "embedding", + [0.1, 0.2, 0.3], + k=10, + distance_fn=l2_distance, + distance_name="score", + ) + sql = ibis.to_sql(expr, dialect="postgres") + assert "L2_DISTANCE" in sql + assert 'AS "score"' in sql + assert '"score" ASC' in sql + assert "LIMIT 10" in sql + + +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 From 44065fcc4f337e4887ced7afadd6ae56bee40b5a Mon Sep 17 00:00:00 2001 From: Rohan Dsouza Date: Wed, 22 Jul 2026 10:27:45 +0530 Subject: [PATCH 2/2] docs: note distance_name/column-name collision in semantic_search Review nit on #40: if table already has a column named distance_name (default "distance"), select raises a duplicate-column error since it ends up in both other_cols and the aliased distance kwarg. --- src/ibis_hotdata/vector.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/ibis_hotdata/vector.py b/src/ibis_hotdata/vector.py index de0d142..1f8dbff 100644 --- a/src/ibis_hotdata/vector.py +++ b/src/ibis_hotdata/vector.py @@ -41,6 +41,9 @@ def semantic_search( Excludes `column` from the result and orders ascending by `distance_fn(column, query_vector)`, aliased as `distance_name`. + + `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()]