diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 1d6ad9897..23bb05605 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -4565,13 +4565,6 @@ class LanceTable(Table): [`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec].""" return LOOP.run(self._table.get_lsm_write_spec()) - def _hybrid_pk_fusion_learned(self) -> bool: - # `self._table` is the async table, which owns the Rust handle. - return self._table._hybrid_pk_fusion_learned() - - def _note_hybrid_pk_fusion(self) -> None: - self._table._note_hybrid_pk_fusion() - def checkpoint_lsm(self) -> None: """Synchronous version of [`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm].""" diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 21ee5e555..4eb7d30e7 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1258,19 +1258,6 @@ fn inject_pk_columns(select: &mut Select, pk_columns: &[String]) { } } -/// Drop `columns` from `batch`, ignoring any that are not present. -fn drop_columns(batch: RecordBatch, columns: &[String]) -> Result { - let keep: Vec = batch - .schema() - .fields() - .iter() - .enumerate() - .filter(|(_, field)| !columns.iter().any(|name| name == field.name())) - .map(|(idx, _)| idx) - .collect(); - Ok(batch.project(&keep)?) -} - impl VectorQuery { fn new(base: Query) -> Self { Self { @@ -1633,7 +1620,13 @@ impl VectorQuery { // has already refused the query if they did. FusionKey::PrimaryKey { injected, .. } => { results = results.drop_column(ROW_ID)?; - results = drop_columns(results, injected)?; + for column in injected { + // A reranker is free to reshape its output, so only drop + // what actually came back. + if results.schema().column_with_name(column).is_some() { + results = results.drop_column(column)?; + } + } } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 3992f8753..15f452f23 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -79,7 +79,7 @@ use std::collections::{HashMap, HashSet}; use std::io::Cursor; use std::pin::Pin; use std::sync::{Arc, Mutex}; -use std::time::{Duration, SystemTime}; +use std::time::{Duration, Instant, SystemTime}; use tokio::sync::RwLock; const REQUEST_TIMEOUT_HEADER: HeaderName = HeaderName::from_static("x-request-timeout-ms"); @@ -91,12 +91,13 @@ const METRIC_TYPE_KEY: &str = "metric_type"; const INDEX_TYPE_KEY: &str = "index_type"; const SCHEMA_CACHE_TTL: Duration = Duration::from_secs(30); const SCHEMA_CACHE_REFRESH_WINDOW: Duration = Duration::from_secs(5); -/// How long a hybrid query remembers that this table refused `_rowid`. Nothing -/// fetches this — it is seeded by the refusal itself — so the TTL only bounds -/// how long the memory outlives a write spec someone removed, and being late to -/// notice costs a base-only read that is still correct. +/// How long a hybrid query remembers that this table refused `_rowid`. +/// +/// Nothing fetches this — the refusal seeds it — so the bound exists only to +/// stop the memory outliving a write spec someone removed through another +/// handle. Being late to notice keeps the query on the primary key, which on a +/// base table is still a valid read; the window is what stops that lasting. const WAL_FUSION_MEMORY_TTL: Duration = Duration::from_secs(30); -const WAL_FUSION_MEMORY_REFRESH_WINDOW: Duration = Duration::from_secs(5); const SCHEMA_SELECTOR_CHANGED: &str = "table selector changed while fetching schema"; fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool { @@ -503,7 +504,7 @@ pub struct RemoteTable { version: Arc>>, location: RwLock>, schema_cache: BackgroundCache, - wal_pk_fusion: BackgroundCache, + wal_pk_fusion: Mutex>, freshness: Arc>, /// The branch this handle is scoped to, or `None` for the main branch. /// Stamped onto every branch-accepting request so reads and writes resolve @@ -672,10 +673,7 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: BackgroundCache::new( - WAL_FUSION_MEMORY_TTL, - WAL_FUSION_MEMORY_REFRESH_WINDOW, - ), + wal_pk_fusion: Mutex::new(None), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, } @@ -709,10 +707,7 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: BackgroundCache::new( - WAL_FUSION_MEMORY_TTL, - WAL_FUSION_MEMORY_REFRESH_WINDOW, - ), + wal_pk_fusion: Mutex::new(None), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch, } @@ -1448,7 +1443,7 @@ impl RemoteTable { } fn invalidate_wal_pk_fusion(&self) { - self.wal_pk_fusion.invalidate(); + *self.wal_pk_fusion.lock().unwrap() = None; } fn handle_error_invalidation(&self, error: &Error) { @@ -1634,10 +1629,7 @@ mod test_utils { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: BackgroundCache::new( - WAL_FUSION_MEMORY_TTL, - WAL_FUSION_MEMORY_REFRESH_WINDOW, - ), + wal_pk_fusion: Mutex::new(None), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, } @@ -1662,10 +1654,7 @@ mod test_utils { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: BackgroundCache::new( - WAL_FUSION_MEMORY_TTL, - WAL_FUSION_MEMORY_REFRESH_WINDOW, - ), + wal_pk_fusion: Mutex::new(None), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, } @@ -1699,10 +1688,7 @@ mod test_utils { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: BackgroundCache::new( - WAL_FUSION_MEMORY_TTL, - WAL_FUSION_MEMORY_REFRESH_WINDOW, - ), + wal_pk_fusion: Mutex::new(None), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, } @@ -3253,7 +3239,6 @@ impl BaseTable for RemoteTable { .json(&body); let (request_id, response) = self.send(request, true).await?; self.check_table_response(&request_id, response).await?; - self.invalidate_wal_pk_fusion(); Ok(()) } @@ -3333,11 +3318,14 @@ impl BaseTable for RemoteTable { } fn hybrid_pk_fusion_learned(&self) -> bool { - self.wal_pk_fusion.try_get().unwrap_or(false) + self.wal_pk_fusion + .lock() + .unwrap() + .is_some_and(|learned| learned.elapsed() < WAL_FUSION_MEMORY_TTL) } fn note_hybrid_pk_fusion(&self) { - self.wal_pk_fusion.seed(true); + *self.wal_pk_fusion.lock().unwrap() = Some(Instant::now()); } async fn tags(&self) -> Result> {