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)