From e40e8d1fa6e6642eb835f9d7b4a45cc22703e079 Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Sat, 27 Jun 2026 16:10:17 -0700 Subject: [PATCH 1/2] perf: lazy + projected payload reads (#116) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reads previously hard-loaded `binary_payload` (and `text_payload` / `embedding`) for every row on every `list` / `search` / `get` (`batch_to_records` required the columns; the scanners didn't project). For multi-modal records that means a metadata `list` or a vector `search` pulled every image's bytes into memory. This adds column projection on the read path plus an on-demand blob fetch: - `ReadProjection { text, binary, embedding }` (default = load all, backward compatible) with `metadata_only()` / `without_binary()` helpers. - `batch_to_records` now tolerates those columns being projected out (omitted fields decode as `None`). - `ContextStore::list_filtered_projected` and `search_filtered_projected` read only the requested columns; the existing `*_with_options` methods delegate with the default (full) projection, so behavior is unchanged. Search always reads the embedding internally to score, then drops it from results if `projection.embedding` is false. - `ContextStore::get_blob(id)` fetches a single record's `binary_payload` on demand via a projected, id-filtered scan. Python: - `Context.list(..., include_binary=True, include_embedding=True)` and `Context.search(..., include_binary=..., include_embedding=...)` — set `include_binary=False` to query metadata/search without materializing media bytes (omitted payloads come back as `None`). - `Context.get_blob(id) -> bytes | None` for on-demand fetch. - The new pyo3 args default to `true`, so existing callers are unaffected. Tests: core (projection drops binary/embedding but keeps metadata; `get_blob` fetches on demand; search projection keeps ranking) and Python (`test_lazy_payload.py`). README gains a multi-modal read example. Independent of #115 (operates on the existing inline payload columns). Follow-up (noted, deferred): REST `include_binary` params + a blob GET route, and Lance `take_blob`/`BlobFile` zero-copy lazy reads for blob-encoded columns. Closes #116 --- README.md | 6 + crates/lance-context-core/src/lib.rs | 2 +- crates/lance-context-core/src/store.rs | 272 +++++++++++++++++++++++-- python/python/lance_context/api.py | 25 ++- python/src/lib.rs | 49 +++-- python/tests/test_lazy_payload.py | 59 ++++++ 6 files changed, 379 insertions(+), 34 deletions(-) create mode 100644 python/tests/test_lazy_payload.py diff --git a/README.md b/README.md index aaa7604..0e7a400 100644 --- a/README.md +++ b/README.md @@ -137,6 +137,12 @@ hits = ctx.search( ) service_context = ctx.related("service://service-a", relation="describes") +# Multi-modal-friendly reads: skip large media bytes for metadata/search +# queries, then fetch a record's bytes on demand. +lean = ctx.list(filters={"content_type": "image/png"}, include_binary=False) +image_bytes = ctx.get_blob(lean[0]["id"]) +hits = ctx.search(query_embedding, include_binary=False) # no bytes pulled into results + # Hybrid retrieval combines lexical recall, vector recall, and existing filters # over the same context records. hybrid_hits = ctx.retrieve( diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index f818aaa..0554a97 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -29,7 +29,7 @@ pub use record::{ }; pub use store::{ CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, DistanceMetric, - IdIndexType, + IdIndexType, ReadProjection, }; // Re-export CompactionMetrics from lance for Python bindings diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 2f64519..1d44b4f 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -277,6 +277,55 @@ struct ExternalIdState { has_non_tombstone: bool, } +/// Which large/optional payload columns to load on a read. +/// +/// Excluding `binary` (and optionally `text` / `embedding`) lets metadata and +/// search queries avoid materializing large media bytes; the omitted fields +/// come back as `None`. Fetch a single record's bytes on demand with +/// [`ContextStore::get_blob`]. Defaults to loading everything (backward +/// compatible). +#[derive(Debug, Clone, Copy)] +pub struct ReadProjection { + pub text: bool, + pub binary: bool, + pub embedding: bool, +} + +impl Default for ReadProjection { + fn default() -> Self { + Self { + text: true, + binary: true, + embedding: true, + } + } +} + +impl ReadProjection { + /// Load only scalar/metadata columns (no text, binary, or embedding). + #[must_use] + pub fn metadata_only() -> Self { + Self { + text: false, + binary: false, + embedding: false, + } + } + + /// Load everything except the (potentially large) `binary_payload`. + #[must_use] + pub fn without_binary() -> Self { + Self { + binary: false, + ..Self::default() + } + } + + fn loads_all(self) -> bool { + self.text && self.binary && self.embedding + } +} + impl ContextStore { /// Open an existing context dataset or create a new one with the project schema. pub async fn open(uri: &str) -> LanceResult { @@ -1164,7 +1213,23 @@ impl ContextStore { filters: Option<&RecordFilters>, options: LifecycleQueryOptions, ) -> LanceResult> { - let scanner = self.lsm_scanner().await?; + self.list_filtered_projected(limit, offset, filters, options, ReadProjection::default()) + .await + } + + /// Like [`Self::list_filtered_with_options`] but with column projection, so + /// large payload columns can be skipped (see [`ReadProjection`]). Omitted + /// payloads come back as `None`; fetch bytes on demand via + /// [`Self::get_blob`]. + pub async fn list_filtered_projected( + &self, + limit: Option, + offset: Option, + filters: Option<&RecordFilters>, + options: LifecycleQueryOptions, + projection: ReadProjection, + ) -> LanceResult> { + let scanner = self.lsm_scanner_projected(projection).await?; let mut stream = scanner.try_into_stream().await?; let mut results = Vec::new(); while let Some(batch) = stream.try_next().await? { @@ -1296,6 +1361,22 @@ impl ContextStore { limit: Option, filters: Option<&RecordFilters>, options: LifecycleQueryOptions, + ) -> LanceResult> { + self.search_filtered_projected(query, limit, filters, options, ReadProjection::default()) + .await + } + + /// Like [`Self::search_filtered_with_options`] but with column projection on + /// the returned records (see [`ReadProjection`]). Embeddings are always read + /// internally to score, then dropped from the results if `projection.embedding` + /// is `false`. + pub async fn search_filtered_projected( + &self, + query: &[f32], + limit: Option, + filters: Option<&RecordFilters>, + options: LifecycleQueryOptions, + projection: ReadProjection, ) -> LanceResult> { validate_query_dimension(query, self.embedding_dim)?; @@ -1304,14 +1385,23 @@ impl ContextStore { return Ok(Vec::new()); } + // Embedding is required to score; force it on for the scan but honor the + // caller's text/binary choices. + let scan_projection = ReadProjection { + embedding: true, + ..projection + }; let mut results: Vec = self - .list_filtered_with_options(None, None, filters, options) + .list_filtered_projected(None, None, filters, options, scan_projection) .await? .into_iter() - .filter_map(|record| { + .filter_map(|mut record| { let distance = self .distance_metric .distance(query, record.embedding.as_ref()?); + if !projection.embedding { + record.embedding = None; + } Some(SearchResult { record, distance }) }) .collect(); @@ -1449,6 +1539,60 @@ impl ContextStore { )) } + /// Top-level column names to read for a projection (drops the excluded + /// payload columns; everything else is always loaded so lifecycle + /// filtering and metadata stay correct). + fn projected_columns(&self, projection: ReadProjection) -> Vec { + self.dataset + .schema() + .fields + .iter() + .map(|field| field.name.clone()) + .filter(|name| { + (projection.text || name != "text_payload") + && (projection.binary || name != "binary_payload") + && (projection.embedding || name != "embedding") + }) + .collect() + } + + /// An LSM scanner that only reads the columns required by `projection`. + async fn lsm_scanner_projected(&self, projection: ReadProjection) -> LanceResult { + let scanner = self.lsm_scanner().await?; + if projection.loads_all() { + return Ok(scanner); + } + let columns = self.projected_columns(projection); + let refs: Vec<&str> = columns.iter().map(String::as_str).collect(); + Ok(scanner.project(&refs)) + } + + /// Fetch a single record's `binary_payload` on demand, without loading it + /// during `list`/`search`. Returns `None` if the record or its binary + /// payload is absent. + pub async fn get_blob(&self, id: &str) -> LanceResult>> { + let filter = format!("id IN ({})", sql_quoted_list(&[id])); + let scanner = self + .lsm_scanner() + .await? + .project(&["id", "binary_payload"]) + .filter(&filter)?; + let mut stream = scanner.try_into_stream().await?; + while let Some(batch) = stream.try_next().await? { + let id_array = column_as::(&batch, "id")?; + let binary_array = column_as_optional::(&batch, "binary_payload"); + for row in 0..batch.num_rows() { + if id_array.value(row) == id { + return Ok(match binary_array { + Some(arr) if !arr.is_null(row) => Some(arr.value(row).to_vec()), + _ => None, + }); + } + } + } + Ok(None) + } + /// Manually trigger compaction to merge small fragments. pub async fn compact( &mut self, @@ -2222,21 +2366,22 @@ fn batch_to_records(batch: &RecordBatch) -> LanceResult> { let supersedes_id_array = column_as_optional::(batch, "supersedes_id"); let superseded_by_id_array = column_as_optional::(batch, "superseded_by_id"); let content_type_array = column_as::(batch, "content_type")?; - let binary_array = column_as::(batch, "binary_payload")?; - let embedding_array = column_as::(batch, "embedding")?; + let binary_array = column_as_optional::(batch, "binary_payload"); + let embedding_array = column_as_optional::(batch, "embedding"); - // Auto-detect whether text_payload is LargeBinary (blob) or LargeUtf8 (default) + // `text_payload` may be projected out, or stored as LargeBinary (blob) or LargeUtf8. + let has_text = batch.schema().field_with_name("text_payload").is_ok(); let text_is_binary = batch .schema() .field_with_name("text_payload") .is_ok_and(|f| f.data_type() == &DataType::LargeBinary); - let text_string_array = if !text_is_binary { + let text_string_array = if has_text && !text_is_binary { Some(column_as::(batch, "text_payload")?) } else { None }; - let text_binary_array = if text_is_binary { + let text_binary_array = if has_text && text_is_binary { Some(column_as::(batch, "text_payload")?) } else { None @@ -2314,32 +2459,30 @@ fn batch_to_records(batch: &RecordBatch) -> LanceResult> { }) }; - let text_payload = if text_is_binary { - let arr = text_binary_array.unwrap(); + let text_payload = if let Some(arr) = text_binary_array { if arr.is_null(row) { None } else { Some(String::from_utf8_lossy(arr.value(row)).to_string()) } - } else { - let arr = text_string_array.unwrap(); + } else if let Some(arr) = text_string_array { if arr.is_null(row) { None } else { Some(arr.value(row).to_string()) } + } else { + None }; - let binary_payload = if binary_array.is_null(row) { - None - } else { - Some(binary_array.value(row).to_vec()) + let binary_payload = match binary_array { + Some(arr) if !arr.is_null(row) => Some(arr.value(row).to_vec()), + _ => None, }; - let embedding = if embedding_array.is_null(row) { - None - } else { - Some(embedding_from_list(embedding_array, row)?) + let embedding = match embedding_array { + Some(arr) if !arr.is_null(row) => Some(embedding_from_list(arr, row)?), + _ => None, }; let role = if role_array.is_null(row) { @@ -4779,4 +4922,93 @@ mod tests { assert_eq!(v1, v2, "ensure_id_index should not recreate existing index"); }); } + + #[test] + fn projection_excludes_binary_but_keeps_metadata() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut record = text_record("img", 0.0); + record.content_type = "image/png".to_string(); + record.binary_payload = Some(vec![1, 2, 3, 4]); + store.add(std::slice::from_ref(&record)).await.unwrap(); + + // Default read still includes the bytes. + let full = store.list(None, None).await.unwrap(); + assert_eq!(full[0].binary_payload.as_deref(), Some(&[1, 2, 3, 4][..])); + assert!(full[0].embedding.is_some()); + + // Projected read drops binary, keeps metadata + embedding. + let projected = store + .list_filtered_projected( + None, + None, + None, + LifecycleQueryOptions::default(), + ReadProjection::without_binary(), + ) + .await + .unwrap(); + assert_eq!(projected.len(), 1); + assert!(projected[0].binary_payload.is_none()); + assert_eq!(projected[0].id, "img"); + assert_eq!(projected[0].content_type, "image/png"); + assert!(projected[0].embedding.is_some()); + + // metadata_only drops embedding too. + let meta = store + .list_filtered_projected( + None, + None, + None, + LifecycleQueryOptions::default(), + ReadProjection::metadata_only(), + ) + .await + .unwrap(); + assert!(meta[0].binary_payload.is_none()); + assert!(meta[0].embedding.is_none()); + assert_eq!(meta[0].id, "img"); + + // get_blob fetches the bytes on demand. + let blob = store.get_blob("img").await.unwrap(); + assert_eq!(blob.as_deref(), Some(&[1, 2, 3, 4][..])); + assert!(store.get_blob("missing").await.unwrap().is_none()); + }); + } + + #[test] + fn search_projection_excludes_binary_keeps_ranking() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_string_lossy().to_string(); + let runtime = tokio::runtime::Runtime::new().unwrap(); + runtime.block_on(async { + let mut store = ContextStore::open(&uri).await.unwrap(); + let mut a = text_record("a", 0.0); + a.binary_payload = Some(vec![9, 9, 9]); + let mut b = text_record("b", 1.0); + b.binary_payload = Some(vec![8, 8, 8]); + store.add(&[a, b]).await.unwrap(); + + let query = make_embedding(0.0); + let results = store + .search_filtered_projected( + &query, + Some(5), + None, + LifecycleQueryOptions::default(), + ReadProjection::without_binary(), + ) + .await + .unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].record.id, "a"); // closest to query pivot 0.0 + assert!(results.iter().all(|r| r.record.binary_payload.is_none())); + // embedding kept by default projection (without_binary keeps embedding) + assert!(results[0].record.embedding.is_some()); + }); + } } diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 3460fb3..5ca1afd 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -1198,6 +1198,8 @@ def search( include_expired: bool = False, include_retired: bool = False, include_relationships: bool = False, + include_binary: bool = True, + include_embedding: bool = True, ) -> list[dict[str, Any]]: provider = getattr(self, "_embedding_provider", None) if isinstance(query, str) and provider is not None: @@ -1210,6 +1212,8 @@ def search( include_expired, include_retired, include_relationships, + include_binary, + include_embedding, ) return [_normalize_search_hit(item) for item in results] @@ -1401,6 +1405,8 @@ def list( *, include_expired: bool = False, include_retired: bool = False, + include_binary: bool = True, + include_embedding: bool = True, ) -> list[dict[str, Any]]: """Return stored entries. @@ -1412,11 +1418,16 @@ def list( created_at range filters, or metadata fields. include_expired: Include records whose ``expires_at`` is in the past. include_retired: Include retired/superseded/revoked records. + include_binary: Load ``binary_payload``. Set ``False`` to avoid + materializing large media bytes for metadata/listing queries; + fetch a record's bytes on demand with :meth:`get_blob`. + include_embedding: Load the embedding vector. Set ``False`` to skip + it when only metadata is needed. Returns: List of entry dicts with keys: id, run_id, role, content_type, text, binary, embedding, created_at, metadata, state_metadata, and - lifecycle metadata. + lifecycle metadata. Omitted payloads come back as ``None``. """ results = self._inner.list( limit, @@ -1424,9 +1435,21 @@ def list( _json_dumps(filters, "filters"), include_expired, include_retired, + include_binary, + include_embedding, ) return [_normalize_record(item) for item in results] + def get_blob(self, id: str) -> bytes | None: + """Fetch a single record's ``binary_payload`` bytes on demand. + + Pairs with ``list(..., include_binary=False)`` / ``search(..., + include_binary=False)``: query metadata cheaply, then pull the media for + the records you actually need. Returns ``None`` if the record or its + binary payload is absent. + """ + return self._inner.get_blob(id) + def get( self, *, id: str | None = None, external_id: str | None = None ) -> dict[str, Any] | None: diff --git a/python/src/lib.rs b/python/src/lib.rs index f341bb1..c905e88 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -23,8 +23,8 @@ use lance_context_core::{ ContextNamespace as RustContextNamespace, ContextRecord, ContextStore, ContextStoreOptions, DistanceMetric, EvalConfig, EvalQuerySet, ExportConfig, ExportTask, GroupBy, IdIndexType, LifecycleQueryOptions, PartitionInfo, PartitionSelector, PartitionSpec, PreferenceForm, - RecordFilters, RecordPatch, Relationship, RetrievalMode, RetrieveResult, SearchResult, - SplitConfig, StateMetadata, LIFECYCLE_ACTIVE, + ReadProjection, RecordFilters, RecordPatch, Relationship, RetrievalMode, RetrieveResult, + SearchResult, SplitConfig, StateMetadata, LIFECYCLE_ACTIVE, }; const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream"; @@ -753,7 +753,7 @@ impl Context { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (query, limit = None, filters_json = None, include_expired = false, include_retired = false, include_relationships = false))] + #[pyo3(signature = (query, limit = None, filters_json = None, include_expired = false, include_retired = false, include_relationships = false, include_binary = true, include_embedding = true))] fn search( &self, py: Python<'_>, @@ -763,17 +763,24 @@ impl Context { include_expired: bool, include_retired: bool, include_relationships: bool, + include_binary: bool, + include_embedding: bool, ) -> PyResult> { let filters = filters_from_json(filters_json)?; let options = LifecycleQueryOptions::new(include_expired, include_retired); + let projection = ReadProjection { + text: true, + binary: include_binary, + embedding: include_embedding, + }; let hits_res = py.allow_threads(|| { - self.runtime - .block_on(self.store.search_filtered_with_options( - &query, - limit, - filters.as_ref(), - options, - )) + self.runtime.block_on(self.store.search_filtered_projected( + &query, + limit, + filters.as_ref(), + options, + projection, + )) }); let hits = hits_res.map_err(to_py_err)?; hits.into_iter() @@ -919,7 +926,8 @@ impl Context { serde_json::to_string(&report).map_err(to_py_err) } - #[pyo3(signature = (limit = None, offset = None, filters_json = None, include_expired = false, include_retired = false))] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (limit = None, offset = None, filters_json = None, include_expired = false, include_retired = false, include_binary = true, include_embedding = true))] fn list( &self, py: Python<'_>, @@ -928,17 +936,25 @@ impl Context { filters_json: Option, include_expired: bool, include_retired: bool, + include_binary: bool, + include_embedding: bool, ) -> PyResult> { let filters = filters_from_json(filters_json)?; let options = LifecycleQueryOptions::new(include_expired, include_retired); + let projection = ReadProjection { + text: true, + binary: include_binary, + embedding: include_embedding, + }; // Release GIL during data retrieval let records = py.allow_threads(|| { self.runtime - .block_on(self.store.list_filtered_with_options( + .block_on(self.store.list_filtered_projected( limit, offset, filters.as_ref(), options, + projection, )) .map_err(to_py_err) })?; @@ -949,6 +965,15 @@ impl Context { .collect() } + fn get_blob(&self, py: Python<'_>, id: &str) -> PyResult> { + let blob = py.allow_threads(|| { + self.runtime + .block_on(self.store.get_blob(id)) + .map_err(to_py_err) + })?; + Ok(blob.map(|bytes| PyBytes::new(py, &bytes).into())) + } + #[pyo3(signature = (target_id, relation = None, limit = None, include_expired = false, include_retired = false))] fn related( &self, diff --git a/python/tests/test_lazy_payload.py b/python/tests/test_lazy_payload.py new file mode 100644 index 0000000..e031668 --- /dev/null +++ b/python/tests/test_lazy_payload.py @@ -0,0 +1,59 @@ +"""Tests for lazy / projected payload reads (issue #116).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context + +IMG = b"\x89PNG\r\n\x1a\n fake image bytes" + + +def test_list_can_exclude_binary_and_fetch_on_demand(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "c.lance"), embedding_dim=4) + ctx.add( + "user", + IMG, + content_type="image/png", + embedding=[1.0, 0.0, 0.0, 0.0], + ) + + # Default read includes the bytes. + full = ctx.list() + assert len(full) == 1 + assert full[0]["binary"] == IMG + assert full[0]["content_type"] == "image/png" + rec_id = full[0]["id"] + + # Excluding binary returns metadata without the bytes. + lean = ctx.list(include_binary=False) + assert lean[0]["binary"] is None + assert lean[0]["content_type"] == "image/png" + assert lean[0]["id"] == rec_id + assert lean[0]["embedding"] is not None # embedding still present + + # metadata-only: drop embedding too. + meta = ctx.list(include_binary=False, include_embedding=False) + assert meta[0]["binary"] is None + assert meta[0]["embedding"] is None + + # Fetch the bytes on demand. + assert ctx.get_blob(rec_id) == IMG + assert ctx.get_blob("does-not-exist") is None + + +def test_search_can_exclude_binary(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "c.lance"), embedding_dim=4) + ctx.add("user", b"img-a", content_type="image/png", embedding=[1.0, 0.0, 0.0, 0.0]) + ctx.add("user", b"img-b", content_type="image/png", embedding=[0.0, 1.0, 0.0, 0.0]) + + hits = ctx.search([1.0, 0.0, 0.0, 0.0], include_binary=False) + assert len(hits) == 2 + assert all(h["binary"] is None for h in hits) + + # Default search still returns the bytes. + full = ctx.search([1.0, 0.0, 0.0, 0.0]) + assert any(h["binary"] is not None for h in full) From f26871900bd19d509cddc09ba6b4fe31a106fc7c Mon Sep 17 00:00:00 2001 From: Allen Cheng Date: Sat, 27 Jun 2026 16:21:17 -0700 Subject: [PATCH 2/2] feat: multi-modal embeddings + cross-modal retrieval (#117) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Embedding providers were text-only (`embed_texts`), so non-text payloads (images/audio) got no embedding unless the caller supplied a vector, and there was no way to retrieve an image with a text query. This adds a multi-modal embedding path (no models bundled — the encoder is user-supplied): - `MultiModalEmbeddingProvider` protocol: extends `EmbeddingProvider` with `embed_media(items: list[tuple[bytes, str]]) -> list[list[float]]`, which embeds `(payload_bytes, content_type)` into the **same** vector space as `embed_texts`. Plus a `supports_media(provider)` helper. - Auto-embed for media: `add` / `upsert` / `add_many` / `upsert_many` now embed `bytes` payloads via `embed_media` when the provider supports it (text payloads still use `embed_texts`). - Cross-modal retrieval falls out of the shared space for free: a text query is embedded with `embed_texts` and matched against image embeddings via the existing vector search — `ctx.search("a photo of a cat")` returns image records. `search` also accepts a `bytes` query (embedded via `embed_media`). Because the two encoders share one space, no core/Rust change is needed — this is entirely in the Python embedding layer. Tests (`test_multimodal_embeddings.py`, deterministic CLIP-style stub): protocol/`supports_media` check, image auto-embed via `embed_media`, cross-modal text→image retrieval, batch auto-embed, and that a text-only provider does not embed images. Stacked on #116 (lazy/projected reads); independent of #115. Note: the 8 failures in the pre-existing `test_embeddings.py` are unrelated mock drift — its hand-written `_DummyInner.add()` accepts 10–18 positional args while the real binding now passes 22 (TypeError at `api.py` `inner.add`, predating this change). The new tests here use a real `Context`, not that mock. Closes #117 --- README.md | 9 ++ python/python/lance_context/__init__.py | 4 + python/python/lance_context/api.py | 59 +++++++--- python/python/lance_context/embeddings.py | 33 +++++- python/tests/test_multimodal_embeddings.py | 125 +++++++++++++++++++++ 5 files changed, 214 insertions(+), 16 deletions(-) create mode 100644 python/tests/test_multimodal_embeddings.py diff --git a/README.md b/README.md index 0e7a400..ef7ee7a 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,15 @@ ctx.add( ) print(ctx.get(external_id="conversation-2026-03-01#turn-1")) ctx.delete(external_id="conversation-2026-03-01#turn-1") + +# Cross-modal retrieval: plug in a multi-modal (e.g. CLIP-style) embedder that +# maps text and media into one shared space. Images are auto-embedded via +# `embed_media`; a text query (embedded via `embed_texts`) then retrieves them. +# lance-context bundles no models — supply your own provider. +clip_ctx = Context.create("multimodal.lance", embedding_dim=512, + embedding_provider=my_clip_provider) # implements MultiModalEmbeddingProvider +clip_ctx.add("user", image_bytes, content_type="image/png", external_id="img-1") +results = clip_ctx.search("a photo of a cat") # text query -> image results assert ctx.get(external_id="conversation-2026-03-01#turn-1") is None # Scoped recall and provenance-oriented metadata diff --git a/python/python/lance_context/__init__.py b/python/python/lance_context/__init__.py index d17869f..da6b8fd 100644 --- a/python/python/lance_context/__init__.py +++ b/python/python/lance_context/__init__.py @@ -8,12 +8,16 @@ RemoteContext, __version__, ) +from .embeddings import ( # pyright: ignore[reportMissingImports] + MultiModalEmbeddingProvider, +) __all__ = [ "AsyncContext", "Context", "ContextNamespace", "EmbeddingProvider", + "MultiModalEmbeddingProvider", "RemoteContext", "__version__", ] diff --git a/python/python/lance_context/api.py b/python/python/lance_context/api.py index 5ca1afd..a24882d 100644 --- a/python/python/lance_context/api.py +++ b/python/python/lance_context/api.py @@ -19,7 +19,7 @@ RemoteContext as _RemoteContext, ) from ._internal import version as _version # pyright: ignore[reportMissingImports] -from .embeddings import EmbeddingProvider, _build_provider +from .embeddings import EmbeddingProvider, _build_provider, supports_media __all__ = [ "AsyncContext", @@ -605,8 +605,11 @@ def add( content_type = data_type payload, resolved_type = _normalize_content(content, content_type) provider = getattr(self, "_embedding_provider", None) - if embedding is None and provider is not None and isinstance(payload, str): - embedding = provider.embed_texts([payload])[0] + if embedding is None and provider is not None: + if isinstance(payload, str): + embedding = provider.embed_texts([payload])[0] + elif isinstance(payload, (bytes, bytearray)) and supports_media(provider): + embedding = provider.embed_media([(bytes(payload), resolved_type)])[0] defaults = getattr(self, "_default_fields", {}) if bot_id is None: bot_id = defaults.get("bot_id") @@ -677,8 +680,11 @@ def upsert( payload, resolved_type = _normalize_content(content, content_type) provider = getattr(self, "_embedding_provider", None) - if embedding is None and provider is not None and isinstance(payload, str): - embedding = provider.embed_texts([payload])[0] + if embedding is None and provider is not None: + if isinstance(payload, str): + embedding = provider.embed_texts([payload])[0] + elif isinstance(payload, (bytes, bytearray)) and supports_media(provider): + embedding = provider.embed_media([(bytes(payload), resolved_type)])[0] defaults = getattr(self, "_default_fields", {}) if bot_id is None: bot_id = defaults.get("bot_id") @@ -1163,21 +1169,39 @@ def iter_rows() -> Iterable[Mapping[str, Any]]: ) def _auto_embed_batch(self, records: list[dict[str, Any]]) -> None: - """Embed text records without an embedding in one provider call.""" + """Embed records without an embedding in one provider call per modality.""" provider = getattr(self, "_embedding_provider", None) if provider is None: return - indices = [ + text_indices = [ i for i, r in enumerate(records) if r.get("embedding") is None and isinstance(r.get("content"), str) ] - if not indices: - return - texts = [records[i]["content"] for i in indices] - vectors = provider.embed_texts(texts) - for i, vec in zip(indices, vectors): - records[i]["embedding"] = vec + if text_indices: + vectors = provider.embed_texts( + [records[i]["content"] for i in text_indices] + ) + for i, vec in zip(text_indices, vectors): + records[i]["embedding"] = vec + if supports_media(provider): + media_indices = [ + i + for i, r in enumerate(records) + if r.get("embedding") is None + and isinstance(r.get("content"), (bytes, bytearray)) + ] + if media_indices: + items = [ + ( + bytes(records[i]["content"]), + records[i].get("data_type") or "application/octet-stream", + ) + for i in media_indices + ] + vectors = provider.embed_media(items) + for i, vec in zip(media_indices, vectors): + records[i]["embedding"] = vec def snapshot(self, label: str | None = None) -> str: return self._inner.snapshot(label) @@ -1202,8 +1226,13 @@ def search( include_embedding: bool = True, ) -> list[dict[str, Any]]: provider = getattr(self, "_embedding_provider", None) - if isinstance(query, str) and provider is not None: - query = provider.embed_texts([query])[0] + if provider is not None: + if isinstance(query, str): + query = provider.embed_texts([query])[0] + elif isinstance(query, (bytes, bytearray)) and supports_media(provider): + query = provider.embed_media( + [(bytes(query), "application/octet-stream")] + )[0] vector = _coerce_vector(query) results = self._inner.search( vector, diff --git a/python/python/lance_context/embeddings.py b/python/python/lance_context/embeddings.py index 8dc2321..e78d082 100644 --- a/python/python/lance_context/embeddings.py +++ b/python/python/lance_context/embeddings.py @@ -2,7 +2,12 @@ from typing import Any, Protocol, runtime_checkable -__all__ = ["EmbeddingProvider", "OpenAIProvider", "SentenceTransformersProvider"] +__all__ = [ + "EmbeddingProvider", + "MultiModalEmbeddingProvider", + "OpenAIProvider", + "SentenceTransformersProvider", +] @runtime_checkable @@ -23,6 +28,32 @@ def embed_texts(self, texts: list[str]) -> list[list[float]]: ... +@runtime_checkable +class MultiModalEmbeddingProvider(EmbeddingProvider, Protocol): + """An embedding provider that also embeds non-text media (images, audio) + into the **same** vector space as :meth:`embed_texts`. + + Sharing one space is what enables **cross-modal retrieval**: an image added + with :meth:`embed_media` and a text query embedded with :meth:`embed_texts` + land in the same space, so ``ctx.search("a photo of a cat")`` can return + image records (no extra wiring — it flows through the normal vector search). + + ``embed_media`` receives ``(payload_bytes, content_type)`` for each item and + must return vectors of the same ``dims`` (and in the same space) as + :meth:`embed_texts`. lance-context bundles no models — supply your own + (e.g. a CLIP-style encoder). + """ + + def embed_media(self, items: list[tuple[bytes, str]]) -> list[list[float]]: + """Return one embedding per ``(payload_bytes, content_type)`` item.""" + ... + + +def supports_media(provider: Any) -> bool: + """Whether ``provider`` can embed non-text media into the shared space.""" + return provider is not None and callable(getattr(provider, "embed_media", None)) + + def _build_provider(config: dict[str, Any]) -> EmbeddingProvider: """Instantiate a built-in provider from a config dict. diff --git a/python/tests/test_multimodal_embeddings.py b/python/tests/test_multimodal_embeddings.py new file mode 100644 index 0000000..b958671 --- /dev/null +++ b/python/tests/test_multimodal_embeddings.py @@ -0,0 +1,125 @@ +"""Tests for multi-modal embeddings + cross-modal retrieval (issue #117). + +Uses a deterministic CLIP-style stub provider (text + media share one vector +space) so no models/deps are needed. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + +from lance_context.api import Context +from lance_context.embeddings import MultiModalEmbeddingProvider, supports_media + + +class StubCLIP: + """Text and media land in the same 4-d space keyed by a concept label.""" + + dims = 4 + + def __init__(self) -> None: + self.text_calls: list[str] = [] + self.media_calls: list[tuple[bytes, str]] = [] + + @staticmethod + def _vec(label: str) -> list[float]: + return { + "cat": [1.0, 0.0, 0.0, 0.0], + "dog": [0.0, 1.0, 0.0, 0.0], + }.get(label, [0.0, 0.0, 1.0, 0.0]) + + def embed_texts(self, texts: list[str]) -> list[list[float]]: + self.text_calls.extend(texts) + return [ + self._vec("cat" if "cat" in t else "dog" if "dog" in t else "other") + for t in texts + ] + + def embed_media(self, items: list[tuple[bytes, str]]) -> list[list[float]]: + self.media_calls.extend(items) + out = [] + for data, _content_type in items: + blob = bytes(data) + label = "cat" if b"cat" in blob else "dog" if b"dog" in blob else "other" + out.append(self._vec(label)) + return out + + +def _ctx(tmp_path: Path) -> Context: + ctx = Context.create(str(tmp_path / "mm.lance"), embedding_dim=4) + ctx._embedding_provider = StubCLIP() # type: ignore[attr-defined] + return ctx + + +def test_protocol_runtime_check() -> None: + provider = StubCLIP() + assert isinstance(provider, MultiModalEmbeddingProvider) + assert supports_media(provider) + + class TextOnly: + dims = 2 + + def embed_texts(self, texts: list[str]) -> list[list[float]]: + return [[0.0, 0.0] for _ in texts] + + assert not supports_media(TextOnly()) + + +def test_image_auto_embedded_via_media(tmp_path: Path) -> None: + ctx = _ctx(tmp_path) + ctx.add("user", b"a cat image", content_type="image/png") + + assert ctx._embedding_provider.media_calls # embed_media was used + rec = ctx.list()[0] + assert rec["embedding"][:2] == [1.0, 0.0] # the shared-space "cat" vector + + +def test_cross_modal_text_query_retrieves_image(tmp_path: Path) -> None: + ctx = _ctx(tmp_path) + ctx.add("user", b"a cat photo", content_type="image/png", external_id="cat") + ctx.add("user", b"a dog photo", content_type="image/png", external_id="dog") + + # Text query embeds via the shared text encoder -> retrieves the image. + hits = ctx.search("a photo of a cat", limit=2) + assert hits[0]["external_id"] == "cat" + assert hits[0]["binary"] == b"a cat photo" + + +def test_batch_auto_embeds_images(tmp_path: Path) -> None: + ctx = _ctx(tmp_path) + ctx.add_many( + [ + { + "role": "user", + "content": b"cat pic", + "content_type": "image/png", + "external_id": "c", + }, + { + "role": "user", + "content": b"dog pic", + "content_type": "image/png", + "external_id": "d", + }, + ] + ) + by_ext = {r["external_id"]: r for r in ctx.list()} + assert by_ext["c"]["embedding"][0] == 1.0 + assert by_ext["d"]["embedding"][1] == 1.0 + + +def test_text_only_provider_does_not_embed_images(tmp_path: Path) -> None: + ctx = Context.create(str(tmp_path / "t.lance"), embedding_dim=2) + + class TextOnly: + dims = 2 + + def embed_texts(self, texts: list[str]) -> list[list[float]]: + return [[1.0, 0.0] for _ in texts] + + ctx._embedding_provider = TextOnly() # type: ignore[attr-defined] + ctx.add("user", b"some bytes") # no embed_media -> not embedded + assert ctx.list()[0]["embedding"] is None