feat: add vector-search helpers (cosine/l2/dot distance UDFs, semantic_search)

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 <distance>(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
 

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 <distance>(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 | ❌ |

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()

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",
+]

src/ibis_hotdata/vector.py

@@ -0,0 +1,55 @@
+"""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`.
+
+    `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))
+    return (
+        table.select(*other_cols, **{distance_name: distance_fn(col, qvec)})
+        .order_by(ibis.asc(distance_name))
+        .limit(k)
+    )

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<float32>"},
+    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