diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index ee58edf3e..4632eb0ea 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -48,22 +48,18 @@ use super::util::parse_server_version; use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id}; mod identifiers; -use identifiers::*; - - +mod namespace_headers; +mod options; mod wire; -use wire::*; +use identifiers::*; +use namespace_headers::*; +pub use options::*; +use wire::*; // the versions of the server that we support // for any new feature that we need to change the SDK behavior, we should bump the server version, // and add a feature flag as method of `ServerVersion` here. -mod options; -pub use options::*; - -mod namespace_headers; -use namespace_headers::*; - pub const DEFAULT_SERVER_VERSION: semver::Version = semver::Version::new(0, 1, 0); #[derive(Debug, Clone)] pub struct ServerVersion(pub semver::Version); @@ -105,7 +101,6 @@ impl ServerVersion { } } - #[derive(Debug)] pub struct RemoteDatabase { client: RestfulLanceDbClient, @@ -120,7 +115,6 @@ pub struct RemoteDatabase { sql_client: Option, } - pub struct RemoteHostOverrides { pub rest: Option, pub sql: Option, @@ -415,8 +409,6 @@ impl From<&CreateTableMode> for &'static str { } } - - /// Bound on `list_jobs` page walking; a warning is logged when the listing /// is truncated at this many pages. const MAX_LIST_JOBS_PAGES: usize = 100; @@ -1453,6 +1445,5 @@ impl Database for RemoteDatabase { } } - #[cfg(test)] mod tests; diff --git a/rust/lancedb/src/remote/db/identifiers.rs b/rust/lancedb/src/remote/db/identifiers.rs index eb6d59644..d01ed3603 100644 --- a/rust/lancedb/src/remote/db/identifiers.rs +++ b/rust/lancedb/src/remote/db/identifiers.rs @@ -19,7 +19,11 @@ pub(super) fn quote_sql_identifier(identifier: &str) -> String { /// `/` that [`RemoteCatalog`] allows. /// /// [`RemoteCatalog`]: super::catalog::RemoteCatalog -pub(super) fn build_object_identifier(what: &str, name: &str, namespace: &[String]) -> Result { +pub(super) fn build_object_identifier( + what: &str, + name: &str, + namespace: &[String], +) -> Result { for segment in namespace { reject_unaddressable_component("namespace segment", segment)?; } diff --git a/rust/lancedb/src/remote/db/tests/catalog.rs b/rust/lancedb/src/remote/db/tests/catalog.rs index 6e5f31657..0829ae749 100644 --- a/rust/lancedb/src/remote/db/tests/catalog.rs +++ b/rust/lancedb/src/remote/db/tests/catalog.rs @@ -348,8 +348,9 @@ async fn test_create_function_async_sends_canonical_request_and_decodes_typed_jo const REQUEST: &str = include_str!( "../../../../tests/fixtures/first_class_functions/v1/remote_function_registration_request.json" ); - const FUNCTION_JOB: &str = - include_str!("../../../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); + const FUNCTION_JOB: &str = include_str!( + "../../../../tests/fixtures/first_class_functions/v1/remote_function_job.json" + ); let mut expected: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); expected.as_object_mut().unwrap().remove("name"); let conn = Connection::new_with_handler(move |request| match request.url().path() { diff --git a/rust/lancedb/src/remote/oauth.rs b/rust/lancedb/src/remote/oauth.rs index 931f033dc..e09f867f1 100644 --- a/rust/lancedb/src/remote/oauth.rs +++ b/rust/lancedb/src/remote/oauth.rs @@ -41,9 +41,9 @@ use crate::error::{Error, Result}; use crate::remote::client::HeaderProvider; mod azure; -use azure::*; - mod interactive; + +use azure::*; use interactive::*; const DEFAULT_REFRESH_BUFFER_SECS: u64 = 300; @@ -969,7 +969,6 @@ impl TokenSource for ClientCredentialsSource { } } - /// Build the token source for a configuration. /// /// Shared by [`OAuthHeaderProvider`] and diff --git a/rust/lancedb/src/remote/oauth/interactive.rs b/rust/lancedb/src/remote/oauth/interactive.rs index c8fe1df50..b32b59cdb 100644 --- a/rust/lancedb/src/remote/oauth/interactive.rs +++ b/rust/lancedb/src/remote/oauth/interactive.rs @@ -317,7 +317,9 @@ impl DeviceCodeSource { }) } - pub(super) async fn request_device_authorization(&self) -> Result { + pub(super) async fn request_device_authorization( + &self, + ) -> Result { let endpoint = self .oidc .get_discovery() diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index edf9c2aee..713a23db2 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -1216,7 +1216,6 @@ fn query_expiration(poll_info: &PollInfo) -> Result) -> Result<()> { +pub(super) fn merge_headers( + destination: &mut HeaderMap, + source: &HashMap, +) -> Result<()> { for (key, value) in source { insert_header(destination, key, value)?; } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index f3ca7593b..6307d6077 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2,7 +2,17 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors pub mod blobs; +mod branches; +mod columns; +mod describe; +mod freshness; +mod indices; pub mod insert; +mod lsm; +mod merge; +mod requests; +mod tags; +mod write; use self::insert::{RemoteWriteExec, WriteOp}; use super::client::RequestResultExt; @@ -84,6 +94,11 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; use tokio::sync::RwLock; +use describe::*; +use freshness::*; +pub use merge::*; +pub use tags::*; + const REQUEST_TIMEOUT_HEADER: HeaderName = HeaderName::from_static("x-request-timeout-ms"); const MIN_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-version"); const MIN_TIMESTAMP_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-timestamp"); @@ -123,261 +138,10 @@ fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool { } } -/// Per-table state driving the freshness headers (`x-lancedb-min-version`, -/// `x-lancedb-min-timestamp`, and `x-lancedb-min-read-version`) sent on table -/// requests. -#[derive(Debug, Default, Clone, Copy)] -struct FreshnessState { - /// Identifies the handle timeline that produced this state. Explicit - /// checkout operations advance the generation so responses from older - /// in-flight requests cannot repopulate the new timeline's constraints. - generation: u64, - /// Exact-version, tag, and snapshot handles must not carry latest-timeline - /// constraints. Their request body already selects the precise version. - pinned: bool, - /// Provides read-your-write within a single handle: writes that return a - /// version update this, and reads send it as `x-lancedb-min-version`. - min_version: Option, - /// Highest committed dataset version advertised by a successful table - /// response on this handle. Later requests send it as - /// `x-lancedb-min-read-version` so a load-balanced query - /// node whose cache is behind this version must refresh before serving, - /// giving monotonic reads across nodes regardless of which one the load - /// balancer routes to. Unlike write result bodies, this is sourced only - /// from the server's committed dataset-version response header or other - /// typed dataset-version fields, so WAL entry ids cannot enter it. - min_read_version: Option, - /// Wall-clock time captured at the last [`BaseTable::checkout_latest`] - /// call. Subsequent reads send - /// `max(baseline, now - read_consistency_interval)` as - /// `x-lancedb-min-timestamp`. - /// - /// Without this, `checkout_latest()` would have no effect on subsequent - /// reads when `read_consistency_interval` is unset (the default): a - /// server-side cache could still serve a snapshot older than the moment - /// the user explicitly asked for "latest". The baseline forces the - /// server to skip any cache entry older than the checkout time, so the - /// `checkout_latest()` signal is preserved across reads on the same - /// handle regardless of the configured consistency interval. - checkout_baseline: Option, -} - -/// Snapshot of the headers that should be attached to a single table request. -#[derive(Debug, Default, Clone, Copy)] -struct FreshnessHeaders { - generation: u64, - min_version: Option, - min_timestamp: Option, - min_read_version: Option, -} - -#[derive(Debug, Default, Clone, Copy)] -struct ReadSnapshot { - version: Option, - freshness_state: FreshnessState, - freshness: FreshnessHeaders, -} - -impl FreshnessHeaders { - fn apply(self, mut request: RequestBuilder) -> RequestBuilder { - if let Some(v) = self.min_version { - request = request.header(MIN_VERSION_HEADER, v.to_string()); - } - if let Some(ts) = self.min_timestamp { - let dt: chrono::DateTime = ts.into(); - request = request.header(MIN_TIMESTAMP_HEADER, dt.to_rfc3339()); - } - if let Some(v) = self.min_read_version { - request = request.header(MIN_READ_VERSION_HEADER, v.to_string()); - } - request - } - - fn observe_version(self, freshness: &Mutex, version: u64) { - track_read_version_for_generation(freshness, self.generation, version); - } - - fn observe_headers( - self, - freshness: &Mutex, - headers: &reqwest::header::HeaderMap, - ) { - if let Some(version) = headers - .get(&VERSION_HEADER) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - { - self.observe_version(freshness, version); - } - } - - fn update_if_current( - self, - freshness: &Mutex, - update: impl FnOnce(&mut FreshnessState), - ) { - let mut state = freshness.lock().unwrap(); - if state.generation == self.generation { - update(&mut state); - } - } - - fn is_current(self, freshness: &Mutex) -> bool { - freshness.lock().unwrap().generation == self.generation - } -} - -fn track_read_version(freshness: &Mutex, version: u64) { - if version == 0 { - return; - } - let mut state = freshness.lock().unwrap(); - state.min_read_version = Some(state.min_read_version.map_or(version, |v| v.max(version))); -} - -fn track_read_version_for_generation( - freshness: &Mutex, - generation: u64, - version: u64, -) { - if version == 0 { - return; - } - let mut state = freshness.lock().unwrap(); - if state.generation == generation { - state.min_read_version = Some(state.min_read_version.map_or(version, |v| v.max(version))); - } -} - -/// A backfill job whose successful wait establishes a read-freshness -/// baseline on the submitting handle, so a later read cannot be served -/// from a cache older than the completed fill. A handle pinned by checkout -/// at completion keeps its time-travel view instead. -struct FreshnessJob { - inner: RemoteJob, - freshness: Arc>, - version: Arc>>, - tracked_result: TrackedJobResult, - freshness_request: FreshnessHeaders, -} - -#[derive(Clone, Copy)] -enum TrackedJobResult { - None, - RefreshColumn, -} - -#[async_trait] -impl crate::job::JobHandle for FreshnessJob { - fn id(&self) -> Option<&str> { - crate::job::JobHandle::id(&self.inner) - } - - async fn status(&self) -> Result { - crate::job::JobHandle::status(&self.inner).await - } - - async fn describe(&self) -> Result { - crate::job::JobHandle::describe(&self.inner).await - } - - async fn events( - &self, - request: crate::job::JobEventsRequest, - ) -> Result> { - crate::job::JobHandle::events(&self.inner, request).await - } - - async fn wait(&self) -> Result { - let result = crate::job::JobHandle::wait(&self.inner).await?; - let version = self.version.read().await; - if version.is_none() { - let result_version = match self.tracked_result { - TrackedJobResult::None => None, - TrackedJobResult::RefreshColumn => result.value().and_then(|value| { - serde_json::from_value::(value.clone()) - .ok() - .map(|result| { - result - .published_version - .map_or(result.source_version, |version| { - version.max(result.source_version) - }) - }) - }), - } - .filter(|version| *version != 0); - if let Some(version) = result_version { - self.freshness_request - .observe_version(&self.freshness, version); - } else { - self.freshness_request - .update_if_current(&self.freshness, |state| { - state.checkout_baseline = Some(SystemTime::now()); - }); - } - } - Ok(result) - } - - async fn cancel(&self) -> Result<()> { - crate::job::JobHandle::cancel(&self.inner).await - } -} - -fn compute_min_timestamp( - state: &FreshnessState, - interval: Option, - now: SystemTime, -) -> Option { - let interval_based = match interval { - None => None, - Some(d) if d.is_zero() => Some(now), - Some(d) => Some(now.checked_sub(d).unwrap_or(now)), - }; - match (interval_based, state.checkout_baseline) { - (None, None) => None, - (Some(t), None) | (None, Some(t)) => Some(t), - (Some(a), Some(b)) => Some(a.max(b)), - } -} - fn quote_sql_identifier(identifier: &str) -> String { format!("\"{}\"", identifier.replace('"', "\"\"")) } -fn freshness_headers_snapshot( - freshness: &Mutex, - interval: Option, -) -> FreshnessHeaders { - freshness_state_snapshot(freshness, interval).1 -} - -fn freshness_state_snapshot( - freshness: &Mutex, - interval: Option, -) -> (FreshnessState, FreshnessHeaders) { - let state = *freshness.lock().unwrap(); - if state.pinned { - return ( - state, - FreshnessHeaders { - generation: state.generation, - ..FreshnessHeaders::default() - }, - ); - } - ( - state, - FreshnessHeaders { - generation: state.generation, - min_version: state.min_version, - min_timestamp: compute_min_timestamp(&state, interval, SystemTime::now()), - min_read_version: state.min_read_version, - }, - ) -} - /// Normalize a branch selector: trim whitespace and treat `""` or `"main"` as /// the (absent) main branch, matching the server's convention. fn normalize_branch(branch: Option) -> Option { @@ -386,109 +150,6 @@ fn normalize_branch(branch: Option) -> Option { .filter(|value| !value.is_empty() && value != "main") } -pub struct RemoteTags<'a, S: HttpSend = Sender> { - inner: &'a RemoteTable, -} - -#[async_trait] -impl Tags for RemoteTags<'_, S> { - async fn list(&self) -> Result> { - let request = self - .inner - .client - .post(&format!("/v1/table/{}/tags/list/", self.inner.identifier)); - let (request_id, response) = self.inner.send_unfenced(request, true).await?; - let response = self - .inner - .check_table_response(&request_id, response) - .await?; - - match response.text().await { - Ok(body) => { - // Explicitly tell serde_json what type we want to deserialize into - let tags_map: HashMap = - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse tags list: {}", e).into(), - request_id, - status_code: None, - })?; - - Ok(tags_map) - } - Err(err) => { - let status_code = err.status(); - Err(Error::Http { - source: Box::new(err), - request_id, - status_code, - }) - } - } - } - - async fn get_version(&self, tag: &str) -> Result { - let request = self.inner.client.post(&format!( - "/v1/table/{}/tags/version/", - self.inner.identifier - )); - self.inner - .resolve_tag_version_with_request(tag, request, false) - .await - } - - async fn create(&mut self, tag: &str, version: u64) -> Result<()> { - let mut body = serde_json::json!({ - "tag": tag, - "version": version - }); - self.inner.apply_branch_body(&mut body); - let request = self - .inner - .client - .post(&format!("/v1/table/{}/tags/create/", self.inner.identifier)) - .json(&body); - - let (request_id, response) = self.inner.send(request, true).await?; - self.inner - .check_table_response(&request_id, response) - .await?; - Ok(()) - } - - async fn delete(&mut self, tag: &str) -> Result<()> { - let request = self - .inner - .client - .post(&format!("/v1/table/{}/tags/delete/", self.inner.identifier)) - .json(&serde_json::json!({ "tag": tag })); - - let (request_id, response) = self.inner.send_unfenced(request, true).await?; - self.inner - .check_table_response(&request_id, response) - .await?; - Ok(()) - } - - async fn update(&mut self, tag: &str, version: u64) -> Result<()> { - let mut body = serde_json::json!({ - "tag": tag, - "version": version - }); - self.inner.apply_branch_body(&mut body); - let request = self - .inner - .client - .post(&format!("/v1/table/{}/tags/update/", self.inner.identifier)) - .json(&body); - - let (request_id, response) = self.inner.send(request, true).await?; - self.inner - .check_table_response(&request_id, response) - .await?; - Ok(()) - } -} - pub struct RemoteTable { client: RestfulLanceDbClient, name: String, @@ -518,1085 +179,6 @@ impl std::fmt::Debug for RemoteTable { } } -impl RemoteTable { - async fn submit_create_index(&self, mut index: IndexBuilder) -> Result> { - self.check_mutable().await?; - let request = self - .client - .post(&format!("/v1/table/{}/create_index/", self.identifier)); - - let column = match index.columns.len() { - 0 => { - return Err(Error::InvalidInput { - message: "No columns specified".into(), - }); - } - 1 => index.columns.pop().unwrap(), - _ => { - return Err(Error::NotSupported { - message: "Indices over multiple columns not yet supported".into(), - }); - } - }; - if matches!( - &index.index, - Index::FTS(params) if params.get_document_granularity().is_list_element() - ) && !self.server_version.support_fts_document_granularity() - { - return Err(Error::NotSupported { - message: "FTS document granularity requires remote server version 0.6.0 or later" - .into(), - }); - } - let schema = self.schema().await?; - let (canonical_column, field) = match &index.index { - Index::FTS(_) => resolve_arrow_fts_field_path(&schema, &column)?, - _ => resolve_arrow_field_path(&schema, &column)?, - }; - let mut body = serde_json::json!({ - "column": canonical_column - }); - - if !index.replace { - body["replace"] = false.into(); - } - - // Add name parameter if provided (for backwards compatibility, only include if Some) - if let Some(ref name) = index.name { - body["name"] = serde_json::Value::String(name.clone()); - } - - // Warn if train=false is specified since it's not meaningful - if !index.train { - log::warn!( - "train=false has no effect remote tables. The index will be created empty and automatically populated in the background." - ); - } - - fn to_json(params: &impl serde::Serialize) -> crate::Result { - serde_json::to_value(params).map_err(|e| Error::InvalidInput { - message: format!("failed to serialize index params {:?}", e), - }) - } - - // Map each Index variant to its wire type name and serializable params. - // Auto is special-cased since it needs schema inspection. - let (index_type_str, params) = match &index.index { - Index::IvfFlat(p) => ("IVF_FLAT", Some(to_json(p)?)), - Index::IvfPq(p) => ("IVF_PQ", Some(to_json(p)?)), - Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)), - Index::IvfHnswSq(p) => ("IVF_HNSW_SQ", Some(to_json(p)?)), - Index::IvfHnswFlat(p) => ("IVF_HNSW_FLAT", Some(to_json(p)?)), - Index::IvfRq(p) => ("IVF_RQ", Some(to_json(p)?)), - Index::BTree(p) => ("BTREE", Some(to_json(p)?)), - Index::Bitmap(p) => ("BITMAP", Some(to_json(p)?)), - Index::LabelList(p) => ("LABEL_LIST", Some(to_json(p)?)), - Index::Fm(p) => ("FM", Some(to_json(p)?)), - Index::ZoneMap(p) => ("ZONEMAP", Some(to_json(p)?)), - Index::NGram(p) => ("NGRAM", Some(to_json(p)?)), - Index::BloomFilter(p) => ("BLOOM_FILTER", Some(to_json(p)?)), - Index::RTree(p) => ("RTREE", Some(to_json(p)?)), - Index::FTS(p) => { - let mut params = to_json(p)?; - if p.get_document_granularity().is_list_element() { - params["document_granularity"] = "list_element".into(); - } - ("FTS", Some(params)) - } - Index::Auto => { - if supported_vector_data_type(field.data_type()) { - body[METRIC_TYPE_KEY] = - serde_json::Value::String(DistanceType::L2.to_string().to_lowercase()); - ("IVF_PQ", None) - } else if supported_btree_data_type(field.data_type()) { - ("BTREE", None) - } else { - return Err(Error::NotSupported { - message: format!( - "there are no indices supported for the field `{}` with the data type {}", - field.name(), - field.data_type() - ), - }); - } - } - _ => { - return Err(Error::NotSupported { - message: "Index type not supported".into(), - }); - } - }; - - body[INDEX_TYPE_KEY] = index_type_str.into(); - if let Some(params) = params { - for (key, value) in params.as_object().expect("params should be a JSON object") { - body[key] = value.clone(); - } - } - self.apply_branch_body(&mut body); - - let request = request.json(&body); - - let (request_id, response) = self.send(request, true).await?; - - let response = self.check_table_response(&request_id, response).await?; - let job_id = response - .text() - .await - .ok() - .and_then(|body| extract_job_id(&body)); - - if let Some(wait_timeout) = index.wait_timeout { - let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); - self.wait_for_index(&[&index_name], wait_timeout).await?; - } - - Ok(job_id) - } - - pub(super) fn new_with_sql_client( - client: RestfulLanceDbClient, - name: String, - namespace: Vec, - identifier: String, - server_version: ServerVersion, - sql_client: Option, - ) -> Self { - Self { - client, - name, - namespace, - identifier, - server_version, - sql_client, - version: Arc::new(RwLock::new(None)), - location: RwLock::new(None), - schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: PkFusionMemory::default(), - freshness: Arc::new(Mutex::new(FreshnessState::default())), - branch: None, - } - } - - /// Seed the schema cache from a `describe` body the caller already fetched. - /// - /// Best effort. `open_table` succeeds today without reading this body, so a - /// body we cannot parse leaves the cache empty and the next schema read - /// fetches it again through the path that reports a real error. - pub(crate) fn seed_schema(&self, describe_body: &str) { - let Ok(description) = serde_json::from_str::(describe_body) else { - return; - }; - self.track_read_version(description.version); - if let Ok(schema) = arrow_schema::Schema::try_from(description.schema) { - self.schema_cache.seed(Arc::new(schema)); - } - } - - /// Return a new handle scoped to `branch`, sharing the client but with fresh - /// caches and version/freshness state (the branch tracks its own latest). - /// Mirrors `NativeTable`'s handle-per-branch model. - fn with_branch(&self, branch: Option) -> Self { - Self { - client: self.client.clone(), - name: self.name.clone(), - namespace: self.namespace.clone(), - identifier: self.identifier.clone(), - server_version: self.server_version.clone(), - sql_client: self.sql_client.clone(), - version: Arc::new(RwLock::new(None)), - location: RwLock::new(None), - schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_pk_fusion: PkFusionMemory::default(), - freshness: Arc::new(Mutex::new(FreshnessState::default())), - branch, - } - } - - /// Stamp the branch onto a request as a `?branch=` query param (used for - /// Arrow-body / query-only ops). `None` (main) leaves the request unchanged, - /// keeping it byte-identical to the non-branch path. - fn apply_branch_query(&self, request: RequestBuilder) -> RequestBuilder { - match &self.branch { - Some(branch) => request.query(&[("branch", branch.as_str())]), - None => request, - } - } - - /// Stamp the branch into a JSON request body under `"branch"` (used for JSON - /// ops). `None` (main) leaves the body unchanged. - fn apply_branch_body(&self, body: &mut serde_json::Value) { - if let Some(branch) = &self.branch { - body["branch"] = serde_json::Value::String(branch.clone()); - } - } - - async fn describe(&self) -> Result { - self.describe_read_snapshot(self.snapshot_read_state().await) - .await - } - - async fn describe_read_snapshot( - &self, - read_snapshot: ReadSnapshot, - ) -> Result { - let request = self - .client - .post(&format!("/v1/table/{}/describe/", self.identifier)); - self.describe_with_request( - request, - read_snapshot.version, - Some(read_snapshot.freshness), - ) - .await - } - - async fn schema_read_snapshot(&self, read_snapshot: ReadSnapshot) -> Result { - if read_snapshot.freshness.is_current(&self.freshness) - && let Some(schema) = self.schema_cache.try_get() - && read_snapshot.freshness.is_current(&self.freshness) - { - return Ok(schema); - } - - let description = self.describe_read_snapshot(read_snapshot).await?; - Ok(Arc::new(description.schema.try_into()?)) - } - - async fn resolve_tag_version_with_request( - &self, - tag: &str, - request: RequestBuilder, - fenced: bool, - ) -> Result { - let request = request.json(&serde_json::json!({ "tag": tag })); - - let (request_id, response) = if fenced { - self.send(request, true).await? - } else { - self.send_unfenced(request, true).await? - }; - let response = self.check_table_response(&request_id, response).await?; - - match response.text().await { - Ok(body) => { - let value: serde_json::Value = - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse tag version: {}", e).into(), - request_id: request_id.clone(), - status_code: None, - })?; - - value - .get("version") - .and_then(|v| v.as_u64()) - .ok_or_else(|| Error::Http { - source: format!("Invalid tag version response: {}", body).into(), - request_id, - status_code: None, - }) - } - Err(err) => { - let status_code = err.status(); - Err(Error::Http { - source: Box::new(err), - request_id, - status_code, - }) - } - } - } - - /// Resolve a tag to its `(branch, version)` coordinate via the `tags/version` - /// endpoint, since the `/branches/create` contract accepts no `from_tag`. - async fn resolve_tag_ref(&self, tag: &str) -> Result<(Option, u64)> { - let request = self - .client - .post(&format!("/v1/table/{}/tags/version/", self.identifier)) - .json(&serde_json::json!({ "tag": tag })); - let (request_id, response) = self.send_unfenced(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - let value: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse tag version: {}", e).into(), - request_id: request_id.clone(), - status_code: None, - })?; - let version = value - .get("version") - .and_then(|v| v.as_u64()) - .ok_or_else(|| Error::Http { - source: format!("Invalid tag version response: {}", body).into(), - request_id, - status_code: None, - })?; - let branch = value - .get("branch") - .and_then(|v| v.as_str()) - .map(String::from); - Ok((normalize_branch(branch), version)) - } - - async fn describe_with_request( - &self, - request: RequestBuilder, - version: Option, - freshness_request: Option, - ) -> Result { - let mut body = serde_json::json!({ "version": version }); - self.apply_branch_body(&mut body); - let request = request.json(&body); - - let (request_id, response) = if let Some(freshness_request) = freshness_request { - self.send_with_freshness(request, true, freshness_request) - .await? - } else { - self.send_unfenced(request, true).await? - }; - - let response = self.check_table_response(&request_id, response).await?; - - match response.text().await { - Ok(body) => { - let description: TableDescription = - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse table description: {}", e).into(), - request_id, - status_code: None, - })?; - if let Some(freshness_request) = freshness_request { - freshness_request.observe_version(&self.freshness, description.version); - } - Ok(description) - } - Err(err) => { - let status_code = err.status(); - Err(Error::Http { - source: Box::new(err), - request_id, - status_code, - }) - } - } - } - - async fn send(&self, req: RequestBuilder, with_retry: bool) -> Result<(String, Response)> { - let freshness_request = self.snapshot_freshness_headers(); - self.send_with_freshness(req, with_retry, freshness_request) - .await - } - - async fn send_with_freshness( - &self, - req: RequestBuilder, - with_retry: bool, - freshness_request: FreshnessHeaders, - ) -> Result<(String, Response)> { - let req = freshness_request.apply(req); - let res = if with_retry { - self.client.send_with_retry(req, None, true).await? - } else { - self.client.send(req).await? - }; - if res.1.status().is_success() { - freshness_request.observe_headers(&self.freshness, res.1.headers()); - } - Ok(res) - } - - async fn send_unfenced( - &self, - req: RequestBuilder, - with_retry: bool, - ) -> Result<(String, Response)> { - if with_retry { - self.client.send_with_retry(req, None, true).await - } else { - self.client.send(req).await - } - } - - pub(super) async fn handle_table_not_found( - table_name: &str, - response: reqwest::Response, - request_id: &str, - ) -> Result { - let status = response.status(); - if status == StatusCode::NOT_FOUND { - let body = response.text().await.ok().unwrap_or_default(); - let request_error = Error::Http { - source: body.into(), - request_id: request_id.into(), - status_code: Some(status), - }; - return Err(Error::TableNotFound { - name: table_name.to_string(), - source: Box::new(request_error), - }); - } - Ok(response) - } - - /// Check if a status code should trigger schema cache invalidation - fn should_invalidate_cache_for_status(status: StatusCode) -> bool { - // Only invalidate for errors that could be schema-related - // Don't invalidate for auth errors (401, 403) or temporary failures (503, 502) - matches!( - status, - StatusCode::BAD_REQUEST // 400 - could be schema mismatch - | StatusCode::NOT_FOUND // 404 - table might have been recreated - | StatusCode::UNPROCESSABLE_ENTITY // 422 - schema validation error - | StatusCode::INTERNAL_SERVER_ERROR // 500 - could be schema issue on server - ) - } - - async fn check_table_response( - &self, - request_id: &str, - response: reqwest::Response, - ) -> Result { - let status = response.status(); - let not_found_result = Self::handle_table_not_found(&self.name, response, request_id).await; - - // Check if we should invalidate cache for 404 errors - if not_found_result.is_err() && Self::should_invalidate_cache_for_status(status) { - self.invalidate_schema_cache(); - } - - let response = not_found_result?; - let result = self.client.check_response(request_id, response).await; - - // Invalidate schema cache on errors that could be schema-related - if result.is_err() && Self::should_invalidate_cache_for_status(status) { - self.invalidate_schema_cache(); - } - - result - } - - async fn read_arrow_response( - &self, - request_id: &str, - response: reqwest::Response, - ) -> Result { - let response = self.check_table_response(request_id, response).await?; - - // The header has to be read before the body, which consumes the response. - let content_type = response - .headers() - .get(CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .map(str::to_owned); - let framing = resolve_arrow_ipc_framing(content_type.as_deref(), request_id)?; - - // Buffer the whole body. File framing keeps its footer at the end, so /query - // cannot decode incrementally. Stream framing could, via - // arrow_ipc::reader::StreamDecoder, but fetch_blobs concatenates every batch - // before returning, so no caller would see data sooner. - let body = response.bytes().await.err_to_http(request_id.into())?; - type IpcBatchIterator = - Box> + Send>; - let (schema, batches): (SchemaRef, IpcBatchIterator) = match framing { - ArrowIpcFraming::Stream => { - let reader = StreamReader::try_new(Cursor::new(body), None)?; - (reader.schema(), Box::new(reader)) - } - ArrowIpcFraming::File => { - let reader = FileReader::try_new(Cursor::new(body), None)?; - (reader.schema(), Box::new(reader)) - } - }; - let stream = futures::stream::iter(batches).map_err(DataFusionError::from); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) - } - - fn apply_query_params( - &self, - body: &mut serde_json::Value, - params: &QueryRequest, - ) -> Result<()> { - params.check_filter()?; - body["prefilter"] = params.prefilter.into(); - // Only forward use_lsm when explicitly set; a server that predates it - // ignores the field and routes as it would by default. - if let Some(use_lsm) = params.use_lsm { - body["use_lsm"] = serde_json::Value::Bool(use_lsm); - } - if let Some(offset) = params.offset { - body["offset"] = serde_json::Value::Number(serde_json::Number::from(offset)); - } - - // Server requires k. - // use isize::MAX as usize to avoid overflow: https://github.com/lancedb/lancedb/issues/2211 - let limit = params.limit.unwrap_or(isize::MAX as usize); - body["k"] = serde_json::Value::Number(serde_json::Number::from(limit)); - - if let Some(filter) = ¶ms.filter { - let filter_sql = match filter { - QueryFilter::Sql(sql) => sql.clone(), - QueryFilter::Datafusion(expr) => expr_to_sql_string(expr)?, - QueryFilter::Substrait(_) => { - return Err(Error::NotSupported { - message: "Substrait filters are not supported for remote queries" - .to_string(), - }); - } - }; - body["filter"] = serde_json::Value::String(filter_sql); - } - - match ¶ms.select { - Select::All => {} - Select::Columns(columns) => { - body["columns"] = serde_json::Value::Array( - columns - .iter() - .map(|s| serde_json::Value::String(s.clone())) - .collect(), - ); - } - Select::Dynamic(pairs) => { - let alias_map = - serde_json::Map::from_iter(pairs.iter().map(|(name, expr)| { - (name.clone(), serde_json::Value::String(expr.clone())) - })); - body["columns"] = alias_map.into(); - } - Select::Expr(pairs) => { - let alias_map: Result> = pairs - .iter() - .map(|(name, expr)| { - expr_to_sql_string(expr) - .map(|sql| (name.clone(), serde_json::Value::String(sql))) - }) - .collect(); - body["columns"] = alias_map?.into(); - } - } - - if params.fast_search { - body["fast_search"] = serde_json::Value::Bool(true); - } - - if params.with_row_id { - body["with_row_id"] = serde_json::Value::Bool(true); - } - - if let Some(full_text_search) = ¶ms.full_text_search { - if full_text_search.wand_factor.is_some() { - return Err(Error::NotSupported { - message: "Wand factor is not yet supported in LanceDB Cloud".into(), - }); - } - - let requires_document_granularity_support = - fts_query_requires_document_granularity_support(&full_text_search.query); - if requires_document_granularity_support - && !self.server_version.support_fts_document_granularity() - { - return Err(Error::NotSupported { - message: - "FTS document granularity requires remote server version 0.6.0 or later" - .into(), - }); - } - - if self.server_version.support_structural_fts() { - body["full_text_query"] = serde_json::json!({ - "query": full_text_search.query.clone(), - }); - } else { - body["full_text_query"] = serde_json::json!({ - "columns": full_text_search.columns().into_iter().collect::>(), - "query": full_text_search.query.query(), - }) - } - } - - if let Some(order_by) = ¶ms.order_by { - body["order_by"] = serde_json::Value::Array( - order_by - .iter() - .map(|o| { - serde_json::json!({ - "column_name": o.column_name, - "ascending": o.ascending, - "nulls_first": o.nulls_first, - }) - }) - .collect(), - ); - } - - Ok(()) - } - - fn apply_vector_query_params( - &self, - mut body: serde_json::Value, - query: &VectorQueryRequest, - ) -> Result> { - self.apply_query_params(&mut body, &query.base)?; - - // Apply general parameters, before we dispatch based on number of query vectors. - if let Some(distance_type) = query.distance_type { - body["distance_type"] = serde_json::json!(distance_type); - } - if let Some(approx_mode) = query.approx_mode { - body["approx_mode"] = serde_json::json!(approx_mode); - } - // In 0.23.1 we migrated from `nprobes` to `minimum_nprobes` and `maximum_nprobes`. - // Old client / new server: since minimum_nprobes is missing, fallback to nprobes - // New client / old server: old server will only see nprobes, make sure to set both - // nprobes and minimum_nprobes - // New client / new server: since minimum_nprobes is present, server can ignore nprobes - body["nprobes"] = query.minimum_nprobes.into(); - body["minimum_nprobes"] = query.minimum_nprobes.into(); - if let Some(maximum_nprobes) = query.maximum_nprobes { - body["maximum_nprobes"] = maximum_nprobes.into(); - } else { - body["maximum_nprobes"] = serde_json::Value::Number(Number::from_u128(0).unwrap()) - } - body["lower_bound"] = query.lower_bound.into(); - body["upper_bound"] = query.upper_bound.into(); - body["ef"] = query.ef.into(); - body["refine_factor"] = query.refine_factor.into(); - if let Some(vector_column) = query.column.as_ref() { - body["vector_column"] = serde_json::Value::String(vector_column.clone()); - } - if !query.use_index { - body["bypass_vector_index"] = serde_json::Value::Bool(true); - } - - fn vector_to_json(vector: &arrow_array::ArrayRef) -> Result { - match vector.data_type() { - DataType::Float32 => { - let array = vector - .as_any() - .downcast_ref::() - .unwrap(); - Ok(serde_json::Value::Array( - array - .values() - .iter() - .map(|v| { - serde_json::Number::from_f64(*v as f64) - .map(serde_json::Value::Number) - .ok_or_else(|| Error::InvalidInput { - message: "query vector must contain only finite values" - .into(), - }) - }) - .collect::>>()?, - )) - } - _ => Err(Error::InvalidInput { - message: "VectorQuery vector must be of type Float32".into(), - }), - } - } - - let bodies = match query.query_vector.len() { - 0 => { - // Server takes empty vector, not null or undefined. - body["vector"] = serde_json::Value::Array(Vec::new()); - vec![body] - } - 1 => { - body["vector"] = vector_to_json(&query.query_vector[0])?; - vec![body] - } - _ => { - if self.server_version.support_multivector() { - let vectors = query - .query_vector - .iter() - .map(vector_to_json) - .collect::>>()?; - body["vector"] = serde_json::Value::Array(vectors); - vec![body] - } else { - // Server does not support multiple vectors in a single query. - // We need to send multiple requests. - let mut bodies = Vec::with_capacity(query.query_vector.len()); - for vector in &query.query_vector { - let mut body = body.clone(); - body["vector"] = vector_to_json(vector)?; - bodies.push(body); - } - bodies - } - } - }; - - Ok(bodies) - } - - async fn create_multipart_write(&self) -> Result { - let request = self.apply_branch_query(self.client.post(&format!( - "/v1/table/{}/multipart_write/create", - self.identifier - ))); - let (request_id, response) = self.send(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse multipart create response: {}", e).into(), - request_id, - status_code: None, - })?; - parsed["upload_id"] - .as_str() - .map(|s| s.to_string()) - .ok_or_else(|| Error::Http { - source: "Missing upload_id in multipart create response".into(), - request_id: String::new(), - status_code: None, - }) - } - - async fn complete_multipart_write(&self, upload_id: &str) -> Result { - let request = self.apply_branch_query( - self.client - .post(&format!( - "/v1/table/{}/multipart_write/complete", - self.identifier - )) - .query(&[("upload_id", upload_id)]), - ); - let (request_id, response) = self.send(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse multipart complete response: {}", e).into(), - request_id, - status_code: None, - })?; - let version = parsed["version"].as_u64().ok_or_else(|| Error::Http { - source: "Missing version in multipart complete response".into(), - request_id: String::new(), - status_code: None, - })?; - Ok(AddResult { version }) - } - - async fn abort_multipart_write(&self, upload_id: &str) -> Result<()> { - let request = self.apply_branch_query( - self.client - .post(&format!( - "/v1/table/{}/multipart_write/abort", - self.identifier - )) - .query(&[("upload_id", upload_id)]), - ); - let (request_id, response) = self.send(request, true).await?; - self.check_table_response(&request_id, response).await?; - Ok(()) - } - - async fn check_mutable(&self) -> Result<()> { - let read_guard = self.version.read().await; - match *read_guard { - None => Ok(()), - Some(version) => Err(Error::NotSupported { - message: format!( - "Cannot mutate table reference fixed at version {}. Call checkout_latest() to get a mutable table reference.", - version - ), - }), - } - } - - async fn snapshot_read_state(&self) -> ReadSnapshot { - let version = self.version.read().await; - let (freshness_state, freshness) = - freshness_state_snapshot(&self.freshness, self.client.read_consistency_interval); - ReadSnapshot { - version: *version, - freshness_state, - freshness, - } - } - - /// Snapshot the freshness headers to attach to a single table request. - /// Computed at call time so that retries reuse the same snapshot. - fn snapshot_freshness_headers(&self) -> FreshnessHeaders { - freshness_headers_snapshot(&self.freshness, self.client.read_consistency_interval) - } - - fn reset_freshness(&self, checkout_baseline: Option, pinned: bool) { - let mut state = self.freshness.lock().unwrap(); - let generation = state.generation.wrapping_add(1); - *state = FreshnessState { - generation, - pinned, - checkout_baseline, - ..FreshnessState::default() - }; - } - - /// Send an LSM operator request with the transport retry layer **off**. - /// - /// Retry policy on these routes belongs to the checkpoint loop, which - /// reads the status and can tell contention from a lost claim. Leaving the - /// transport layer on would re-ask on its own schedule first, and surface - /// an `Error::Retry` whose status the loop would then have to unwrap. - async fn send_lsm_route(&self, request: RequestBuilder) -> Result<(String, reqwest::Response)> { - let (request_id, response) = self.send(request, false).await?; - let response = self.check_table_response(&request_id, response).await?; - Ok((request_id, response)) - } - - /// Record a version returned by a write so subsequent reads can request at - /// least that version via `x-lancedb-min-version`. A returned `0` from a - /// backward-compatible old server is ignored. - fn track_write_version(&self, freshness_request: FreshnessHeaders, version: u64) { - if version == 0 { - return; - } - freshness_request.update_if_current(&self.freshness, |state| { - state.min_version = Some(state.min_version.map_or(version, |v| v.max(version))); - }); - } - - /// Record a committed dataset version observed in a table response so - /// subsequent requests ask for at least this version via - /// `x-lancedb-min-read-version`, - /// giving monotonic reads across load-balanced query nodes. A returned `0` - /// (or absent header from an old server) is ignored. - fn track_read_version(&self, version: u64) { - track_read_version(&self.freshness, version); - } - - async fn execute_query( - &self, - query: &AnyQuery, - options: &QueryExecutionOptions, - ) -> Result>>> { - let mut request = self - .client - .post(&format!("/v1/table/{}/query/", self.identifier)); - - if let Some(timeout) = options.timeout { - // Also send to server, so it can abort the query if it takes too long. - // (If it doesn't fit into u64, it's not worth sending anyways.) - if let Ok(timeout_ms) = u64::try_from(timeout.as_millis()) { - request = request.header(REQUEST_TIMEOUT_HEADER, timeout_ms); - } - } - - let read_snapshot = self.snapshot_read_state().await; - let query_bodies = self.prepare_query_bodies(query, read_snapshot.version)?; - let requests: Vec = query_bodies - .into_iter() - .map(|body| request.try_clone().unwrap().json(&body)) - .collect(); - - let futures = requests.into_iter().map(|req| async move { - let (request_id, response) = self - .send_with_freshness(req, true, read_snapshot.freshness) - .await?; - self.read_arrow_response(&request_id, response).await - }); - let streams = futures::future::try_join_all(futures); - - if let Some(timeout) = options.timeout { - let timeout_future = tokio::time::sleep(timeout); - tokio::pin!(timeout_future); - tokio::pin!(streams); - tokio::select! { - _ = &mut timeout_future => { - Err(Error::Other { - message: format!("Query timeout after {} ms", timeout.as_millis()), - source: None, - }) - } - result = &mut streams => { - Ok(result?) - } - } - } else { - Ok(streams.await?) - } - } - - fn prepare_query_bodies( - &self, - query: &AnyQuery, - version: Option, - ) -> Result> { - let query = query.canonicalized()?; - let mut base_body = serde_json::json!({ "version": version }); - self.apply_branch_body(&mut base_body); - - match &query { - AnyQuery::Query(query) => { - let mut body = base_body.clone(); - self.apply_query_params(&mut body, query)?; - // Empty vector can be passed if no vector search is performed. - body["vector"] = serde_json::Value::Array(Vec::new()); - Ok(vec![body]) - } - AnyQuery::VectorQuery(query) => self.apply_vector_query_params(base_body, query), - } - } - - fn invalidate_schema_cache(&self) { - self.schema_cache.invalidate(); - } - - fn handle_error_invalidation(&self, error: &Error) { - let status_code = match error { - Error::Http { status_code, .. } => *status_code, - Error::Retry { status_code, .. } => *status_code, - _ => None, - }; - if let Some(status_code) = status_code - && Self::should_invalidate_cache_for_status(status_code) - { - self.invalidate_schema_cache(); - } - } -} - -#[derive(Deserialize)] -struct TableDescription { - version: u64, - schema: JsonSchema, - location: Option, -} - -/// How a response body frames its Arrow IPC payload. `/query` answers with file framing -/// and `fetch_blobs` with stream framing, so the reader is picked per response. -enum ArrowIpcFraming { - File, - Stream, -} - -/// A response with no `Content-Type` uses file framing, preserving this helper's -/// behavior before `fetch_blobs` introduced stream responses. -fn resolve_arrow_ipc_framing( - content_type: Option<&str>, - request_id: &str, -) -> Result { - let Some(media_type) = content_type.map(base_media_type) else { - return Ok(ArrowIpcFraming::File); - }; - if media_type.eq_ignore_ascii_case(ARROW_STREAM_CONTENT_TYPE) { - return Ok(ArrowIpcFraming::Stream); - } - if media_type.eq_ignore_ascii_case(ARROW_FILE_CONTENT_TYPE) { - return Ok(ArrowIpcFraming::File); - } - Err(Error::Http { - source: format!( - "Expected an Arrow IPC response with Content-Type '{ARROW_STREAM_CONTENT_TYPE}' \ - or '{ARROW_FILE_CONTENT_TYPE}', got '{media_type}'" - ) - .into(), - request_id: request_id.into(), - status_code: None, - }) -} - -/// An Arrow IPC stream carrying `schema` and no batches. -/// -/// The body of an all-null `add_columns`: the server reads the fields to -/// add straight off the schema message, so nothing but the field list is -/// worth sending. -fn write_ipc_schema(schema: &arrow_schema::Schema) -> Result> { - let mut body = Vec::new(); - { - let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, schema)?; - writer.finish()?; - } - Ok(body) -} - -/// Strip media-type parameters before matching against the Arrow content types. -fn base_media_type(content_type: &str) -> &str { - match content_type.split_once(';') { - Some((media_type, _parameters)) => media_type.trim(), - None => content_type.trim(), - } -} - -/// Extract an Error from Arc, reconstructing if the Arc is shared. -/// This is needed because `Shared` futures cache results internally, so -/// `Arc::try_unwrap` typically fails. -fn unwrap_shared_error(arc: Arc) -> Error { - match Arc::try_unwrap(arc) { - Ok(err) => err, - Err(arc) => match &*arc { - Error::TableNotFound { name, source } => Error::TableNotFound { - name: name.clone(), - source: source.to_string().into(), - }, - _ => Error::Runtime { - message: arc.to_string(), - }, - }, - } -} - -async fn fetch_schema( - client: &RestfulLanceDbClient, - identifier: &str, - table_name: &str, - read_snapshot: ReadSnapshot, - branch: Option, - freshness: Arc>, -) -> Result { - let mut body = serde_json::json!({ "version": read_snapshot.version }); - if let Some(branch) = &branch { - body["branch"] = serde_json::Value::String(branch.clone()); - } - let freshness_headers = read_snapshot.freshness; - let request = freshness_headers - .apply(client.post(&format!("/v1/table/{}/describe/", identifier))) - .json(&body); - - let (request_id, response) = client.send_with_retry(request, None, true).await?; - - if response.status() == StatusCode::NOT_FOUND { - let body = response.text().await.ok().unwrap_or_default(); - return Err(Error::TableNotFound { - name: table_name.to_string(), - source: Box::new(Error::Http { - source: body.into(), - request_id, - status_code: Some(StatusCode::NOT_FOUND), - }), - }); - } - - let response = client.check_response(&request_id, response).await?; - freshness_headers.observe_headers(&freshness, response.headers()); - let body = response.text().await.map_err(|e| { - let status_code = e.status(); - Error::Http { - source: Box::new(e), - request_id: request_id.clone(), - status_code, - } - })?; - - let description: TableDescription = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse table description: {}", e).into(), - request_id, - status_code: None, - })?; - freshness_headers.observe_version(&freshness, description.version); - if !freshness_headers.is_current(&freshness) { - return Err(Error::Runtime { - message: SCHEMA_SELECTOR_CHANGED.to_string(), - }); - } - - let arrow_schema: arrow_schema::Schema = description.schema.try_into()?; - Ok(Arc::new(arrow_schema)) -} - impl std::fmt::Display for RemoteTable { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "RemoteTable({})", self.identifier) @@ -1698,403 +280,6 @@ mod test_utils { } } -impl RemoteTable { - fn is_retryable_write_error(&self, err: &Error) -> bool { - match err { - Error::Http { - source, - status_code, - .. - } => { - // Don't retry read errors (is_body/is_decode): the - // server may have committed the write already, and - // without an idempotency key we'd duplicate data. - source - .downcast_ref::() - .is_some_and(|e| e.is_connect()) - || status_code.is_some_and(|s| self.client.retry_config.statuses.contains(&s)) - } - // send_with_retry exhausted its internal retries on a retryable - // status. The outer loop can still retry the whole operation with - // a fresh session. - Error::Retry { status_code, .. } => { - status_code.is_some_and(|s| self.client.retry_config.statuses.contains(&s)) - } - _ => false, - } - } - - async fn add_single_partition(&self, output: PreprocessingOutput) -> Result { - use crate::remote::retry::RetryCounter; - - let _guard = output.tracker.as_ref().map(|t| t.track_task()); - let freshness_request = self.snapshot_freshness_headers(); - - let mut insert: Arc = Arc::new( - RemoteWriteExec::new( - self.name.clone(), - self.identifier.clone(), - self.client.clone(), - output.plan, - WriteOp::Insert { - overwrite: output.overwrite, - }, - output.tracker.clone(), - self.branch.clone(), - ) - .with_freshness( - self.freshness.clone(), - self.client.read_consistency_interval, - ), - ); - - let mut retry_counter = - RetryCounter::new(&self.client.retry_config, uuid::Uuid::new_v4().to_string()); - - loop { - let stream = execute_plan(insert.clone(), Default::default())?; - let result: Result> = stream.try_collect().await.map_err(Error::from); - - match result { - Ok(_) => { - let add_result = (insert.as_ref() as &dyn std::any::Any) - .downcast_ref::>() - .and_then(|i| i.add_result()) - .unwrap_or(AddResult { version: 0 }); - - if output.overwrite { - self.invalidate_schema_cache(); - } - self.track_write_version(freshness_request, add_result.version); - - return Ok(add_result); - } - Err(err) if output.rescannable && self.is_retryable_write_error(&err) => { - retry_counter.increment_from_error(err)?; - tokio::time::sleep(retry_counter.next_sleep_time()).await; - insert = insert.reset_state()?; - continue; - } - Err(err) => return Err(err), - } - } - } - - async fn add_multipart( - &self, - output: PreprocessingOutput, - num_partitions: usize, - ) -> Result { - use crate::remote::retry::RetryCounter; - - let mut retry_counter = - RetryCounter::new(&self.client.retry_config, uuid::Uuid::new_v4().to_string()); - - loop { - let freshness_request = self.snapshot_freshness_headers(); - let upload_id = self.create_multipart_write().await?; - - let result = self - .execute_multipart_inserts(&upload_id, &output, num_partitions) - .await; - - match result { - Ok(()) => match self.complete_multipart_write(&upload_id).await { - Ok(result) => { - if output.overwrite { - self.invalidate_schema_cache(); - } - self.track_write_version(freshness_request, result.version); - return Ok(result); - } - Err(e) => { - if let Err(abort_err) = self.abort_multipart_write(&upload_id).await { - log::warn!( - "Failed to abort multipart write {}: {}", - upload_id, - abort_err - ); - } - if output.rescannable && self.is_retryable_write_error(&e) { - retry_counter.increment_from_error(e)?; - tokio::time::sleep(retry_counter.next_sleep_time()).await; - continue; - } - return Err(e); - } - }, - Err(e) => { - if let Err(abort_err) = self.abort_multipart_write(&upload_id).await { - log::warn!( - "Failed to abort multipart write {}: {}", - upload_id, - abort_err - ); - } - if output.rescannable && self.is_retryable_write_error(&e) { - retry_counter.increment_from_error(e)?; - tokio::time::sleep(retry_counter.next_sleep_time()).await; - continue; - } - return Err(e); - } - } - } - } - - async fn execute_multipart_inserts( - &self, - upload_id: &str, - output: &PreprocessingOutput, - num_partitions: usize, - ) -> Result<()> { - debug_assert!( - output.rescannable, - "multipart inserts require rescannable input for retry support" - ); - - let plan = Arc::new( - datafusion_physical_plan::repartition::RepartitionExec::try_new( - output.plan.clone(), - datafusion_physical_plan::Partitioning::RoundRobinBatch(num_partitions), - )?, - ) as Arc; - - let insert = Arc::new( - RemoteWriteExec::new_multipart( - self.name.clone(), - self.identifier.clone(), - self.client.clone(), - plan, - output.overwrite, - upload_id.to_string(), - output.tracker.clone(), - self.branch.clone(), - self.client.max_bytes_per_request(), - self.client.max_request_duration(), - ) - .with_freshness( - self.freshness.clone(), - self.client.read_consistency_interval, - ), - ); - - let task_ctx = Arc::new(datafusion_execution::TaskContext::default()); - let tracker = output.tracker.clone(); - let mut join_set = tokio::task::JoinSet::new(); - for partition in 0..num_partitions { - let exec = insert.clone(); - let ctx = task_ctx.clone(); - let tracker = tracker.clone(); - join_set.spawn(async move { - let _guard = tracker.as_ref().map(|t| t.track_task()); - let mut stream = exec - .execute(partition, ctx) - .map_err(|e| -> Error { e.into() })?; - while let Some(batch) = stream.next().await { - batch.map_err(|e| -> Error { e.into() })?; - } - Ok::<_, Error>(()) - }); - } - - // JoinSet aborts all remaining tasks when dropped, so if we return - // early on error the orphaned tasks are automatically cancelled. - while let Some(result) = join_set.join_next().await { - result.map_err(|e| Error::Runtime { - message: format!("Insert task panicked: {}", e), - })??; - } - - Ok(()) - } -} - -/// Deserialize an index's `created_at` field. -/// -/// The server returns this as an RFC 3339 string (e.g. `"2026-06-18T21:37:36.637Z"`), -/// but older deployments sent a unix timestamp in milliseconds. Accept both so the -/// client works against any server version. -fn deserialize_created_at<'de, D>( - deserializer: D, -) -> std::result::Result>, D::Error> -where - D: serde::Deserializer<'de>, -{ - use serde::de::Error as _; - - #[derive(Deserialize)] - #[serde(untagged)] - enum CreatedAt { - Rfc3339(String), - Millis(i64), - } - - match Option::::deserialize(deserializer)? { - None => Ok(None), - Some(CreatedAt::Rfc3339(s)) => DateTime::parse_from_rfc3339(&s) - .map(|dt| Some(dt.with_timezone(&Utc))) - .map_err(D::Error::custom), - Some(CreatedAt::Millis(ms)) => Ok(DateTime::from_timestamp_millis(ms)), - } -} - -impl RemoteTable { - async fn index_stats_read_snapshot( - &self, - index_name: &str, - read_snapshot: ReadSnapshot, - ) -> Result> { - let encoded_name = urlencoding::encode(index_name); - let mut body = serde_json::json!({ "version": read_snapshot.version }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!( - "/v1/table/{}/index/{encoded_name}/stats/", - self.identifier - )) - .json(&body); - - let (request_id, response) = self - .send_with_freshness(request, true, read_snapshot.freshness) - .await?; - if response.status() == StatusCode::NOT_FOUND { - return Ok(None); - } - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - let stats = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse index statistics: {}", e).into(), - request_id, - status_code: None, - })?; - Ok(Some(stats)) - } - - /// Parse the response from `/index/list/` into `IndexConfig` entries. - /// - /// When the server returns `index_type` inline, all enriched fields are - /// used directly and no further requests are made. When `index_type` is - /// absent (legacy servers), a `/index/{name}/stats/` call is made for each - /// index to retrieve the type. - async fn parse_index_list_response( - &self, - body: &str, - request_id: &str, - schema: &SchemaRef, - read_snapshot: ReadSnapshot, - ) -> Result> { - use crate::index::IndexType; - - #[derive(Deserialize)] - struct ListIndicesResponse { - indexes: Vec, - } - - #[derive(Deserialize)] - struct IndexListEntry { - index_name: String, - columns: Vec, - // Present on enriched responses; absent on legacy servers. - // Used as the sentinel to decide whether to skip the stats call. - index_type: Option, - index_uuid: Option, - #[serde(default, deserialize_with = "deserialize_created_at")] - created_at: Option>, - num_indexed_rows: Option, - num_unindexed_rows: Option, - size_bytes: Option, - num_segments: Option, - index_version: Option, - index_details: Option, - type_url: Option, - } - - let response: ListIndicesResponse = - serde_json::from_str(body).map_err(|err| Error::Http { - source: format!( - "Failed to parse list_indices response: {}, body: {}", - err, body - ) - .into(), - request_id: request_id.to_string(), - status_code: None, - })?; - - let mut futures = Vec::with_capacity(response.indexes.len()); - for entry in response.indexes { - let columns = entry - .columns - .iter() - .map(|column| { - resolve_arrow_field_path(schema, column) - .map(|(canonical_column, _)| canonical_column) - }) - .collect::>>()?; - - let future = async move { - if let Some(index_type) = entry.index_type { - // Enriched response: all fields available, no stats call needed. - Ok(Some(IndexConfig { - name: entry.index_name, - index_type, - columns, - index_uuid: entry.index_uuid, - type_url: entry.type_url, - created_at: entry.created_at, - num_indexed_rows: entry.num_indexed_rows, - num_unindexed_rows: entry.num_unindexed_rows, - size_bytes: entry.size_bytes, - num_segments: entry.num_segments, - index_version: entry.index_version, - index_details: entry.index_details, - })) - } else { - // Legacy response: fetch index type via stats endpoint. - match self - .index_stats_read_snapshot(&entry.index_name, read_snapshot) - .await - { - Ok(Some(stats)) => Ok(Some(IndexConfig { - name: entry.index_name, - index_type: stats.index_type, - columns, - index_uuid: None, - type_url: None, - created_at: None, - num_indexed_rows: None, - num_unindexed_rows: None, - size_bytes: None, - num_segments: None, - index_version: None, - index_details: None, - })), - Ok(None) => Ok(None), // Index deleted since we listed it. - Err(e) => Err(e), - } - } - }; - futures.push(future); - } - - let results = futures::future::try_join_all(futures).await?; - let mut indices: Vec = results.into_iter().flatten().collect(); - let lance_schema = lance_core::datatypes::Schema::try_from(schema.as_ref())?; - for index in &mut indices { - if index.index_type == IndexType::FTS { - // The wire format uses physical paths for schema resolution. Match - // native tables by exposing list-transparent paths to callers. - for column in &mut index.columns { - let field_id = lance_schema.field_id(column)?; - *column = public_fts_field_path_by_id(&lance_schema, field_id)?; - } - } - } - Ok(indices) - } -} - #[async_trait] impl BaseTable for RemoteTable { fn as_any(&self) -> &dyn std::any::Any { @@ -2404,77 +589,7 @@ impl BaseTable for RemoteTable { name: &str, from: lance::dataset::refs::Ref, ) -> Result> { - use lance::dataset::refs::Ref; - - if name.trim().is_empty() { - return Err(Error::InvalidInput { - message: "branch name must be a non-empty string".into(), - }); - } - - // Translate the source ref into the `from_branch` / `from_version` the - // `/branches/create` contract accepts (it has no `from_tag`). - let (from_branch, from_version) = match from { - Ref::Version(branch, version) => (normalize_branch(branch), version), - Ref::VersionNumber(version) => (normalize_branch(self.branch.clone()), Some(version)), - Ref::Tag(tag) => { - let (branch, version) = self.resolve_tag_ref(&tag).await?; - (branch, Some(version)) - } - }; - - let mut body = serde_json::json!({ "name": name }); - if let Some(from_branch) = &from_branch { - body["from_branch"] = serde_json::Value::String(from_branch.clone()); - } - if let Some(from_version) = from_version { - body["from_version"] = serde_json::json!(from_version); - } - - let request = self - .client - .post(&format!("/v1/table/{}/branches/create/", self.identifier)) - .json(&body); - - // Send without retry so the expected 409 (branch already exists) is - // surfaced as a response we can map, rather than being retried. - let (request_id, response) = self.send_unfenced(request, false).await?; - match response.status() { - StatusCode::CONFLICT => { - return Err(Error::TableAlreadyExists { - name: format!("{} (branch: {})", self.name, name), - }); - } - StatusCode::BAD_REQUEST => { - let body = response.text().await.unwrap_or_default(); - return Err(Error::InvalidInput { - message: format!("invalid create_branch request: {}", body), - }); - } - StatusCode::NOT_FOUND => { - // 404 covers both a missing table and a missing source ref; name - // the source coordinate so the error isn't misattributed to the table. - let body = response.text().await.unwrap_or_default(); - let source_desc = match (&from_branch, from_version) { - (Some(b), Some(v)) => format!(" (source: branch '{b}' version {v})"), - (Some(b), None) => format!(" (source: branch '{b}')"), - (None, Some(v)) => format!(" (source: version {v})"), - (None, None) => String::new(), - }; - return Err(Error::TableNotFound { - name: format!("{}{}", self.name, source_desc), - source: Box::new(Error::Http { - source: body.into(), - request_id, - status_code: Some(StatusCode::NOT_FOUND), - }), - }); - } - _ => {} - } - self.check_table_response(&request_id, response).await?; - - Ok(Arc::new(self.with_branch(Some(name.to_string())))) + self.create_branch_impl(name, from).await } async fn checkout_branch(&self, name: &str) -> Result> { @@ -2497,143 +612,19 @@ impl BaseTable for RemoteTable { } async fn list_branches(&self) -> Result> { - use lance::dataset::refs::BranchContents; - - let request = self - .client - .post(&format!("/v1/table/{}/branches/list/", self.identifier)); - let (request_id, response) = self.send_unfenced(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - #[derive(Deserialize)] - struct ListBranchesResponse { - branches: HashMap, - } - - let parsed: ListBranchesResponse = - serde_json::from_str(&body).map_err(|err| Error::Http { - source: format!( - "Failed to parse list_branches response: {}, body: {}", - err, body - ) - .into(), - request_id, - status_code: None, - })?; - - Ok(parsed.branches) + self.list_branches_impl().await } async fn delete_branch(&self, name: &str) -> Result<()> { - if name.trim().is_empty() { - return Err(Error::InvalidInput { - message: "branch name must be a non-empty string".into(), - }); - } - let request = self - .client - .post(&format!("/v1/table/{}/branches/delete/", self.identifier)) - .json(&serde_json::json!({ "name": name })); - let (request_id, response) = self.send(request, true).await?; - if response.status() == StatusCode::NOT_FOUND { - return Err(Error::TableNotFound { - name: format!("{} (branch: {})", self.name, name), - source: format!("branch '{}' does not exist", name).into(), - }); - } - self.check_table_response(&request_id, response).await?; - Ok(()) + self.delete_branch_impl(name).await } async fn diff_branch(&self, from_branch: &str) -> Result { - if from_branch.trim().is_empty() { - return Err(Error::InvalidInput { - message: "Branch name cannot be empty.".into(), - }); - } - let request = self - .client - .post(&format!("/v1/table/{}/branches/diff/", self.identifier)) - .json(&serde_json::json!({ "from_branch": from_branch })); - let (request_id, response) = self.send_unfenced(request, true).await?; - if response.status() == StatusCode::NOT_FOUND { - return Err(Error::TableNotFound { - name: format!("{} (branch: {})", self.name, from_branch), - source: format!("branch '{}' does not exist", from_branch).into(), - }); - } - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - serde_json::from_str(&body).map_err(|err| Error::Http { - source: format!( - "Failed to parse diff_branch response: {}, body: {}", - err, body - ) - .into(), - request_id, - status_code: None, - }) + self.diff_branch_impl(from_branch).await } async fn cherry_pick(&self, from_branch: &str, dry_run: bool) -> Result { - if from_branch.trim().is_empty() { - return Err(Error::InvalidInput { - message: "Branch name cannot be empty.".into(), - }); - } - let read_snapshot = self.snapshot_read_state().await; - let target_freshness = if self.branch.is_none() && read_snapshot.version.is_none() { - Some(read_snapshot.freshness) - } else { - None - }; - let request = self - .client - .post(&format!( - "/v1/table/{}/branches/cherry_pick/", - self.identifier - )) - .json(&serde_json::json!({ - "from_branch": from_branch, - "dry_run": dry_run, - })); - // No retry. HTTP 409 is CherryPickStatus::Failed with a body, not a transport error. - let (request_id, response) = self.send_unfenced(request, false).await?; - let status = response.status(); - if status == StatusCode::NOT_FOUND { - return Err(Error::TableNotFound { - name: format!("{} (branch: {})", self.name, from_branch), - source: format!("branch '{}' does not exist", from_branch).into(), - }); - } - // 200 and 409 both carry CherryPickResult. - if status != StatusCode::OK && status != StatusCode::CONFLICT { - let body = response.text().await.unwrap_or_default(); - return Err(Error::Http { - source: format!("unexpected status {status} from cherry_pick: {body}").into(), - request_id, - status_code: Some(status), - }); - } - let body = response.text().await.err_to_http(request_id.clone())?; - let result: CherryPickResult = serde_json::from_str(&body).map_err(|err| Error::Http { - source: format!( - "Failed to parse cherry_pick response: {}, body: {}", - err, body - ) - .into(), - request_id, - status_code: Some(status), - })?; - if !dry_run - && status == StatusCode::OK - && result.status == crate::table::CherryPickStatus::CherryPicked - && let (Some(freshness), Some(version)) = (target_freshness, result.main_version_after) - { - freshness.observe_version(&self.freshness, version); - } - Ok(result) + self.cherry_pick_impl(from_branch, dry_run).await } fn current_branch(&self) -> Option { @@ -3160,139 +1151,19 @@ impl BaseTable for RemoteTable { } async fn get_lsm_stats(&self, include_generation_rows: bool) -> Result> { - // Read-semantics POST, like `get_lsm_write_spec`. - let request = self - .client - .post(&format!("/v1/table/{}/get_lsm_stats/", self.identifier)) - .json(&serde_json::json!({ - "include_generation_rows": include_generation_rows, - })); - let (request_id, response) = self.send_lsm_route(request).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - let parsed: GetLsmStatsResponse = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse get_lsm_stats response: {e}").into(), - request_id, - status_code: None, - })?; - // `null` — and only — when the table has no LSM write path. - Ok(parsed.lsm_stats) + self.get_lsm_stats_impl(include_generation_rows).await } async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> { - self.check_mutable().await?; - - // Map the spec onto the server's request DTO. `sharding` is internally - // tagged on `mode` to mirror sophon's `Sharding` enum. A null - // `maintained_indexes` asks the server to resolve every maintainable - // index at HEAD; a list is verbatim, an empty one meaning none. - let sharding = match &spec { - LsmWriteSpec::Bucket { - column, - num_buckets, - .. - } => serde_json::json!({ - "mode": "bucket", - "column": column, - "num_buckets": num_buckets, - }), - LsmWriteSpec::Identity { column, .. } => serde_json::json!({ - "mode": "identity", - "column": column, - }), - LsmWriteSpec::Unsharded { .. } => serde_json::json!({ "mode": "unsharded" }), - }; - let body = serde_json::json!({ - "sharding": sharding, - "maintained_indexes": spec.maintained_indexes(), - "writer_config_defaults": spec.writer_config_defaults(), - }); - - let request = self - .client - .post(&format!( - "/v1/table/{}/set_lsm_write_spec/", - self.identifier - )) - .json(&body); - let (request_id, response) = self.send(request, true).await?; - self.check_table_response(&request_id, response).await?; - Ok(()) + self.set_lsm_write_spec_impl(spec).await } async fn unset_lsm_write_spec(&self) -> Result<()> { - self.check_mutable().await?; - let request = self.client.post(&format!( - "/v1/table/{}/unset_lsm_write_spec/", - self.identifier - )); - let (request_id, response) = self.send(request, true).await?; - self.check_table_response(&request_id, response).await?; - self.wal_pk_fusion.forget(); // WAL-PK-FUSION: delete. - Ok(()) + self.unset_lsm_write_spec_impl().await } async fn get_lsm_write_spec(&self) -> Result> { - // Read counterpart to set/unset, resolved server-side against HEAD. The - // server reads the spec from the `__lance_mem_wal` system index (shard - // column mapped from its Lance field id against the current schema) and - // re-encodes it into the same sophon-owned shape the set endpoint - // accepts — no lance/lancedb types cross the wire. `lsm_write_spec` is - // null when the LSM write path is not enabled for the table. - let request = self.client.post(&format!( - "/v1/table/{}/get_lsm_write_spec/", - self.identifier - )); - let (request_id, response) = self.send(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - // Mirror of sophon's `Sharding` (internally tagged on `mode`) and - // `LsmWriteSpecBody` / `GetLsmWriteSpecResponse`. - #[derive(Deserialize)] - #[serde(tag = "mode", rename_all = "snake_case")] - enum Sharding { - Unsharded, - Bucket { column: String, num_buckets: u32 }, - Identity { column: String }, - } - #[derive(Deserialize)] - struct LsmWriteSpecBody { - sharding: Sharding, - /// `null` selects every index the table has; `[]` selects none. - #[serde(default)] - maintained_indexes: Option>, - #[serde(default)] - writer_config_defaults: std::collections::HashMap, - } - #[derive(Deserialize)] - struct GetLsmWriteSpecResponse { - lsm_write_spec: Option, - } - - let parsed: GetLsmWriteSpecResponse = - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse get_lsm_write_spec response: {}", e).into(), - request_id, - status_code: None, - })?; - - let Some(body) = parsed.lsm_write_spec else { - // The LSM write path is not enabled for this table. - return Ok(None); - }; - - let spec = match body.sharding { - Sharding::Bucket { - column, - num_buckets, - } => LsmWriteSpec::bucket(column, num_buckets), - Sharding::Identity { column } => LsmWriteSpec::identity(column), - Sharding::Unsharded => LsmWriteSpec::unsharded(), - } - .with_maintained_indexes(body.maintained_indexes) - .with_writer_config_defaults(body.writer_config_defaults); - - Ok(Some(spec)) + self.get_lsm_write_spec_impl().await } // WAL-PK-FUSION: delete both. @@ -3339,124 +1210,11 @@ impl BaseTable for RemoteTable { transforms: NewColumnTransform, _read_columns: Option>, ) -> Result { - self.check_mutable().await?; - crate::table::computed_columns::ensure_not_function_bound( - self.schema().await?.as_ref(), - "schema evolution", - crate::table::schema_evolution::new_column_names(&transforms), - )?; - let path = format!("/v1/table/{}/add_columns/", self.identifier); - let request = match transforms { - NewColumnTransform::SqlExpressions(expressions) => { - let body = expressions - .into_iter() - .map(|(name, expression)| { - serde_json::json!({ - "name": name, - "expression": expression, - }) - }) - .collect::>(); - let mut body = serde_json::json!({ "new_columns": body }); - self.apply_branch_body(&mut body); - self.client.post(&path).json(&body) - } - // Every field becomes a column that reads null for existing - // rows. The schema goes over the wire as Arrow IPC rather than - // a JSON description of the types: that is what the server - // takes, and it is the only encoding that round-trips a field - // whole, including decimal precision and a timestamp's unit and - // timezone. With no JSON envelope the branch rides the query - // string, as it does for the other binary-bodied endpoints. - NewColumnTransform::AllNulls(schema) => { - let body = write_ipc_schema(schema.as_ref())?; - self.apply_branch_query( - self.client - .post(&path) - .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) - .body(body), - ) - } - _ => { - return Err(Error::NotSupported { - message: "Only SQL expressions and all-null column schemas are supported \ - for adding columns" - .into(), - }); - } - }; - - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - if body.trim().is_empty() { - // Backward compatible with old servers - return Ok(AddColumnsResult { version: 0 }); - } - - let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse add_columns response: {}", e).into(), - request_id, - status_code: None, - })?; - - self.invalidate_schema_cache(); - self.track_write_version(freshness_request, result.version); - - Ok(result) + self.add_columns_impl(transforms, _read_columns).await } async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result { - self.check_mutable().await?; - crate::table::computed_columns::ensure_not_function_bound( - self.schema().await?.as_ref(), - "schema evolution", - columns.iter().map(|(name, _)| name), - )?; - // The server plans the declaration against its table schema, including - // Blob v2 semantics inherited by a direct field projection. - let entries = columns - .iter() - .map( - |(name, expression)| lance_namespace::models::AddColumnsEntry { - name: name.clone(), - computed: Some(Some(expression.clone())), - ..Default::default() - }, - ) - .collect::>(); - let mut body = serde_json::json!({ "new_columns": entries }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!("/v1/table/{}/add_columns/", self.identifier)) - .json(&body); - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - if body.trim().is_empty() { - // Backward compatible with old servers - return Ok(AddColumnsResult { version: 0 }); - } - - let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse add_columns response: {}", e).into(), - request_id, - status_code: None, - })?; - - self.invalidate_schema_cache(); - self.track_write_version(freshness_request, result.version); - - Ok(result) + self.add_computed_columns_impl(columns).await } async fn add_function_columns( @@ -3464,64 +1222,8 @@ impl BaseTable for RemoteTable { application: &crate::function::FunctionApplication, output_name: Option<&str>, ) -> Result { - self.check_mutable().await?; - let schema = self.schema().await?; - let plan = crate::table::computed_columns::plan_function_application( - schema.as_ref(), - application, - output_name, - )?; - let new_columns = plan - .outputs - .iter() - .map(|output| { - serde_json::json!({ - "name": output.output_name, - "all_null": true, - }) - }) - .collect::>(); - let mut body = serde_json::json!({ - "new_columns": new_columns, - "function": { - "application": plan.application, - "binding_metadata_version": plan.binding_metadata_version, - "input_bindings": plan.input_bindings, - "input_schema": plan.input_schema, - "output_schema": plan.output_schema, - "outputs": plan.outputs, - }, - }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!("/v1/table/{}/add_columns/", self.identifier)) - .json(&body); - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - // A Function declaration can return 404 for the Function rather than the table. - let response = self - .client - .check_response(&request_id, response) + self.add_function_columns_impl(application, output_name) .await - .inspect_err(|error| self.handle_error_invalidation(error))?; - let body = response.text().await.err_to_http(request_id.clone())?; - - if body.trim().is_empty() { - return Ok(AddColumnsResult { version: 0 }); - } - - let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse add Function columns response: {e}").into(), - request_id, - status_code: None, - })?; - - self.invalidate_schema_cache(); - self.track_write_version(freshness_request, result.version); - Ok(result) } async fn refresh_column(&self, _column: &str) -> Result { @@ -3538,185 +1240,29 @@ impl BaseTable for RemoteTable { &self, column: &str, ) -> Result> { - self.check_mutable().await?; - let mut body = serde_json::json!({ "column": column }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!("/v1/table/{}/backfill_column", self.identifier)) - .json(&body); - let (request_id, response) = self.send(request, true).await?; - // Preserve dependency errors: a deleted bound Function also returns 404. - let response = self - .client - .check_response(&request_id, response) - .await - .inspect_err(|error| self.handle_error_invalidation(error))?; - let body = response.text().await.err_to_http(request_id.clone())?; - - #[derive(serde::Deserialize)] - struct BackfillResponse { - job_id: String, - } - let response: BackfillResponse = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse backfill_column response: {}", e).into(), - request_id, - status_code: None, - })?; - - Ok(Job::new_typed(Box::new(FreshnessJob { - inner: RemoteJob::new(self.client.clone(), response.job_id), - freshness: self.freshness.clone(), - version: self.version.clone(), - tracked_result: TrackedJobResult::RefreshColumn, - freshness_request: self.snapshot_freshness_headers(), - }))) + self.refresh_column_async_impl(column).await } async fn function_errors( &self, request: &crate::function::FunctionErrorsRequest, ) -> Result { - let mut body = serde_json::json!({}); - if let Some(job_id) = &request.job_id { - body["job_id"] = serde_json::json!(job_id); - } - if let Some(column) = &request.column { - body["column"] = serde_json::json!(column); - } - if let Some(limit) = request.limit { - body["limit"] = serde_json::json!(limit); - } - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!("/v1/table/{}/errors", self.identifier)) - .json(&body); - let (request_id, response) = self.send(request, true).await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse errors response: {}", e).into(), - request_id, - status_code: None, - }) + self.function_errors_impl(request).await } async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result { - self.check_mutable().await?; - let body = alterations - .iter() - .map(|alteration| { - let mut value = serde_json::json!({ - "path": alteration.path, - }); - if let Some(rename) = &alteration.rename { - value["rename"] = serde_json::Value::String(rename.clone()); - } - if let Some(data_type) = &alteration.data_type { - let json_data_type = JsonDataType::try_from(data_type).unwrap(); - let json_data_type = serde_json::to_value(&json_data_type).unwrap(); - value["data_type"] = json_data_type; - } - if let Some(nullable) = &alteration.nullable { - value["nullable"] = serde_json::Value::Bool(*nullable); - } - value - }) - .collect::>(); - let mut body = serde_json::json!({ "alterations": body }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!("/v1/table/{}/alter_columns/", self.identifier)) - .json(&body); - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - if body.trim().is_empty() { - // Backward compatible with old servers - return Ok(AlterColumnsResult { version: 0 }); - } - - let result: AlterColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse alter_columns response: {}", e).into(), - request_id, - status_code: None, - })?; - - self.invalidate_schema_cache(); - self.track_write_version(freshness_request, result.version); - - Ok(result) + self.alter_columns_impl(alterations).await } async fn update_field_metadata( &self, updates: &[FieldMetadataUpdate], ) -> Result { - self.check_mutable().await?; - let mut body = serde_json::json!({ "updates": updates }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!( - "/v1/table/{}/update_field_metadata/", - self.identifier - )) - .json(&body); - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - let result: UpdateFieldMetadataResult = - serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse update_field_metadata response: {}", e).into(), - request_id, - status_code: None, - })?; - - self.invalidate_schema_cache(); - self.track_write_version(freshness_request, result.version); - Ok(result) + self.update_field_metadata_impl(updates).await } async fn drop_columns(&self, columns: &[&str]) -> Result { - self.check_mutable().await?; - let mut body = serde_json::json!({ "columns": columns }); - self.apply_branch_body(&mut body); - let request = self - .client - .post(&format!("/v1/table/{}/drop_columns/", self.identifier)) - .json(&body); - let freshness_request = self.snapshot_freshness_headers(); - let (request_id, response) = self - .send_with_freshness(request, true, freshness_request) - .await?; - let response = self.check_table_response(&request_id, response).await?; - let body = response.text().await.err_to_http(request_id.clone())?; - - if body.trim().is_empty() { - // Backward compatible with old servers - return Ok(DropColumnsResult { version: 0 }); - } - - let result: DropColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { - source: format!("Failed to parse drop_columns response: {}", e).into(), - request_id, - status_code: None, - })?; - - self.invalidate_schema_cache(); - self.track_write_version(freshness_request, result.version); - - Ok(result) + self.drop_columns_impl(columns).await } async fn list_indices(&self) -> Result> { @@ -3873,96 +1419,5 @@ impl BaseTable for RemoteTable { } } -#[derive(Serialize, Clone, Debug)] -pub struct MergeInsertRequest { - // Sent as one repeated `on` query parameter per column, which is how the - // namespace spec encodes an array-valued `on`. serde_urlencoded (which - // reqwest's `query()` uses) cannot serialize a sequence nested in a struct, - // so this field is emitted separately by [`Self::on_query_params`]. - #[serde(skip_serializing)] - on: Vec, - when_matched_update_all: bool, - when_matched_update_all_filt: Option, - when_not_matched_insert_all: bool, - when_not_matched_by_source_delete: bool, - when_not_matched_by_source_delete_filt: Option, - // For backwards compatibility, only serialize use_index when it's false - // (the default is true) - #[serde(skip_serializing_if = "is_true")] - use_index: bool, - // Only serialize use_lsm when explicitly set (Some); a server that predates - // it ignores the field and routes as it would by default. - #[serde(skip_serializing_if = "Option::is_none")] - use_lsm: Option, -} - -impl MergeInsertRequest { - /// The `on` columns as repeated query parameters: `?on=a&on=b`. - /// - /// A single column serializes to `?on=a`, exactly what clients sent before - /// `on` became a list, so a server that predates composite keys sees no - /// change from a single-column caller. - pub(crate) fn on_query_params(&self) -> Vec<(&str, &str)> { - self.on.iter().map(|col| ("on", col.as_str())).collect() - } -} - -fn is_true(b: &bool) -> bool { - *b -} - -impl TryFrom for MergeInsertRequest { - type Error = Error; - - fn try_from(value: MergeInsertBuilder) -> Result { - if value.on.is_empty() { - return Err(Error::InvalidInput { - message: "MergeInsertBuilder missing required 'on' field".into(), - }); - } - // The server rejects a repeated column with a 400; catching it here - // names the offending column and costs no round trip. - let mut seen = HashSet::with_capacity(value.on.len()); - if let Some(dup) = value.on.iter().find(|col| !seen.insert(*col)) { - return Err(Error::InvalidInput { - message: format!("MergeInsertBuilder 'on' column '{dup}' is repeated"), - }); - } - - let when_matched_update_all_filt = match value.when_matched_update_all_filt { - Some(MergeFilter::Sql(sql)) => Some(sql), - Some(MergeFilter::Expr(_)) => { - return Err(Error::NotSupported { - message: "DataFusion expressions are not supported on remote tables".into(), - }); - } - None => None, - }; - - let when_not_matched_by_source_delete_filt = - match value.when_not_matched_by_source_delete_filt { - Some(MergeFilter::Sql(sql)) => Some(sql), - Some(MergeFilter::Expr(_)) => { - return Err(Error::NotSupported { - message: "DataFusion expressions are not supported on remote tables".into(), - }); - } - None => None, - }; - - Ok(Self { - on: value.on, - when_matched_update_all: value.when_matched_update_all, - when_matched_update_all_filt, - when_not_matched_insert_all: value.when_not_matched_insert_all, - when_not_matched_by_source_delete: value.when_not_matched_by_source_delete, - when_not_matched_by_source_delete_filt, - // Only serialize use_index when it's false for backwards compatibility - use_index: value.use_index, - use_lsm: value.use_lsm, - }) - } -} - #[cfg(test)] mod tests; diff --git a/rust/lancedb/src/remote/table/branches.rs b/rust/lancedb/src/remote/table/branches.rs new file mode 100644 index 000000000..2c8ba5be9 --- /dev/null +++ b/rust/lancedb/src/remote/table/branches.rs @@ -0,0 +1,230 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +impl RemoteTable { + pub(super) async fn create_branch_impl( + &self, + name: &str, + from: lance::dataset::refs::Ref, + ) -> Result> { + use lance::dataset::refs::Ref; + + if name.trim().is_empty() { + return Err(Error::InvalidInput { + message: "branch name must be a non-empty string".into(), + }); + } + + // Translate the source ref into the `from_branch` / `from_version` the + // `/branches/create` contract accepts (it has no `from_tag`). + let (from_branch, from_version) = match from { + Ref::Version(branch, version) => (normalize_branch(branch), version), + Ref::VersionNumber(version) => (normalize_branch(self.branch.clone()), Some(version)), + Ref::Tag(tag) => { + let (branch, version) = self.resolve_tag_ref(&tag).await?; + (branch, Some(version)) + } + }; + + let mut body = serde_json::json!({ "name": name }); + if let Some(from_branch) = &from_branch { + body["from_branch"] = serde_json::Value::String(from_branch.clone()); + } + if let Some(from_version) = from_version { + body["from_version"] = serde_json::json!(from_version); + } + + let request = self + .client + .post(&format!("/v1/table/{}/branches/create/", self.identifier)) + .json(&body); + + // Send without retry so the expected 409 (branch already exists) is + // surfaced as a response we can map, rather than being retried. + let (request_id, response) = self.send_unfenced(request, false).await?; + match response.status() { + StatusCode::CONFLICT => { + return Err(Error::TableAlreadyExists { + name: format!("{} (branch: {})", self.name, name), + }); + } + StatusCode::BAD_REQUEST => { + let body = response.text().await.unwrap_or_default(); + return Err(Error::InvalidInput { + message: format!("invalid create_branch request: {}", body), + }); + } + StatusCode::NOT_FOUND => { + // 404 covers both a missing table and a missing source ref; name + // the source coordinate so the error isn't misattributed to the table. + let body = response.text().await.unwrap_or_default(); + let source_desc = match (&from_branch, from_version) { + (Some(b), Some(v)) => format!(" (source: branch '{b}' version {v})"), + (Some(b), None) => format!(" (source: branch '{b}')"), + (None, Some(v)) => format!(" (source: version {v})"), + (None, None) => String::new(), + }; + return Err(Error::TableNotFound { + name: format!("{}{}", self.name, source_desc), + source: Box::new(Error::Http { + source: body.into(), + request_id, + status_code: Some(StatusCode::NOT_FOUND), + }), + }); + } + _ => {} + } + self.check_table_response(&request_id, response).await?; + + Ok(Arc::new(self.with_branch(Some(name.to_string())))) + } + + pub(super) async fn list_branches_impl( + &self, + ) -> Result> { + use lance::dataset::refs::BranchContents; + + let request = self + .client + .post(&format!("/v1/table/{}/branches/list/", self.identifier)); + let (request_id, response) = self.send_unfenced(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + #[derive(Deserialize)] + struct ListBranchesResponse { + branches: HashMap, + } + + let parsed: ListBranchesResponse = + serde_json::from_str(&body).map_err(|err| Error::Http { + source: format!( + "Failed to parse list_branches response: {}, body: {}", + err, body + ) + .into(), + request_id, + status_code: None, + })?; + + Ok(parsed.branches) + } + + pub(super) async fn delete_branch_impl(&self, name: &str) -> Result<()> { + if name.trim().is_empty() { + return Err(Error::InvalidInput { + message: "branch name must be a non-empty string".into(), + }); + } + let request = self + .client + .post(&format!("/v1/table/{}/branches/delete/", self.identifier)) + .json(&serde_json::json!({ "name": name })); + let (request_id, response) = self.send(request, true).await?; + if response.status() == StatusCode::NOT_FOUND { + return Err(Error::TableNotFound { + name: format!("{} (branch: {})", self.name, name), + source: format!("branch '{}' does not exist", name).into(), + }); + } + self.check_table_response(&request_id, response).await?; + Ok(()) + } + + pub(super) async fn diff_branch_impl(&self, from_branch: &str) -> Result { + if from_branch.trim().is_empty() { + return Err(Error::InvalidInput { + message: "Branch name cannot be empty.".into(), + }); + } + let request = self + .client + .post(&format!("/v1/table/{}/branches/diff/", self.identifier)) + .json(&serde_json::json!({ "from_branch": from_branch })); + let (request_id, response) = self.send_unfenced(request, true).await?; + if response.status() == StatusCode::NOT_FOUND { + return Err(Error::TableNotFound { + name: format!("{} (branch: {})", self.name, from_branch), + source: format!("branch '{}' does not exist", from_branch).into(), + }); + } + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + serde_json::from_str(&body).map_err(|err| Error::Http { + source: format!( + "Failed to parse diff_branch response: {}, body: {}", + err, body + ) + .into(), + request_id, + status_code: None, + }) + } + + pub(super) async fn cherry_pick_impl( + &self, + from_branch: &str, + dry_run: bool, + ) -> Result { + if from_branch.trim().is_empty() { + return Err(Error::InvalidInput { + message: "Branch name cannot be empty.".into(), + }); + } + let read_snapshot = self.snapshot_read_state().await; + let target_freshness = if self.branch.is_none() && read_snapshot.version.is_none() { + Some(read_snapshot.freshness) + } else { + None + }; + let request = self + .client + .post(&format!( + "/v1/table/{}/branches/cherry_pick/", + self.identifier + )) + .json(&serde_json::json!({ + "from_branch": from_branch, + "dry_run": dry_run, + })); + // No retry. HTTP 409 is CherryPickStatus::Failed with a body, not a transport error. + let (request_id, response) = self.send_unfenced(request, false).await?; + let status = response.status(); + if status == StatusCode::NOT_FOUND { + return Err(Error::TableNotFound { + name: format!("{} (branch: {})", self.name, from_branch), + source: format!("branch '{}' does not exist", from_branch).into(), + }); + } + // 200 and 409 both carry CherryPickResult. + if status != StatusCode::OK && status != StatusCode::CONFLICT { + let body = response.text().await.unwrap_or_default(); + return Err(Error::Http { + source: format!("unexpected status {status} from cherry_pick: {body}").into(), + request_id, + status_code: Some(status), + }); + } + let body = response.text().await.err_to_http(request_id.clone())?; + let result: CherryPickResult = serde_json::from_str(&body).map_err(|err| Error::Http { + source: format!( + "Failed to parse cherry_pick response: {}, body: {}", + err, body + ) + .into(), + request_id, + status_code: Some(status), + })?; + if !dry_run + && status == StatusCode::OK + && result.status == crate::table::CherryPickStatus::CherryPicked + && let (Some(freshness), Some(version)) = (target_freshness, result.main_version_after) + { + freshness.observe_version(&self.freshness, version); + } + Ok(result) + } +} diff --git a/rust/lancedb/src/remote/table/columns.rs b/rust/lancedb/src/remote/table/columns.rs new file mode 100644 index 000000000..4b56b5bd6 --- /dev/null +++ b/rust/lancedb/src/remote/table/columns.rs @@ -0,0 +1,387 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +impl RemoteTable { + pub(super) async fn add_columns_impl( + &self, + transforms: NewColumnTransform, + _read_columns: Option>, + ) -> Result { + self.check_mutable().await?; + crate::table::computed_columns::ensure_not_function_bound( + self.schema().await?.as_ref(), + "schema evolution", + crate::table::schema_evolution::new_column_names(&transforms), + )?; + let path = format!("/v1/table/{}/add_columns/", self.identifier); + let request = match transforms { + NewColumnTransform::SqlExpressions(expressions) => { + let body = expressions + .into_iter() + .map(|(name, expression)| { + serde_json::json!({ + "name": name, + "expression": expression, + }) + }) + .collect::>(); + let mut body = serde_json::json!({ "new_columns": body }); + self.apply_branch_body(&mut body); + self.client.post(&path).json(&body) + } + // Every field becomes a column that reads null for existing + // rows. The schema goes over the wire as Arrow IPC rather than + // a JSON description of the types: that is what the server + // takes, and it is the only encoding that round-trips a field + // whole, including decimal precision and a timestamp's unit and + // timezone. With no JSON envelope the branch rides the query + // string, as it does for the other binary-bodied endpoints. + NewColumnTransform::AllNulls(schema) => { + let body = write_ipc_schema(schema.as_ref())?; + self.apply_branch_query( + self.client + .post(&path) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .body(body), + ) + } + _ => { + return Err(Error::NotSupported { + message: "Only SQL expressions and all-null column schemas are supported \ + for adding columns" + .into(), + }); + } + }; + + let freshness_request = self.snapshot_freshness_headers(); + let (request_id, response) = self + .send_with_freshness(request, true, freshness_request) + .await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + if body.trim().is_empty() { + // Backward compatible with old servers + return Ok(AddColumnsResult { version: 0 }); + } + + let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse add_columns response: {}", e).into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(freshness_request, result.version); + + Ok(result) + } + + pub(super) async fn add_computed_columns_impl( + &self, + columns: &[(String, String)], + ) -> Result { + self.check_mutable().await?; + crate::table::computed_columns::ensure_not_function_bound( + self.schema().await?.as_ref(), + "schema evolution", + columns.iter().map(|(name, _)| name), + )?; + // The server plans the declaration against its table schema, including + // Blob v2 semantics inherited by a direct field projection. + let entries = columns + .iter() + .map( + |(name, expression)| lance_namespace::models::AddColumnsEntry { + name: name.clone(), + computed: Some(Some(expression.clone())), + ..Default::default() + }, + ) + .collect::>(); + let mut body = serde_json::json!({ "new_columns": entries }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/add_columns/", self.identifier)) + .json(&body); + let freshness_request = self.snapshot_freshness_headers(); + let (request_id, response) = self + .send_with_freshness(request, true, freshness_request) + .await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + if body.trim().is_empty() { + // Backward compatible with old servers + return Ok(AddColumnsResult { version: 0 }); + } + + let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse add_columns response: {}", e).into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(freshness_request, result.version); + + Ok(result) + } + + pub(super) async fn add_function_columns_impl( + &self, + application: &crate::function::FunctionApplication, + output_name: Option<&str>, + ) -> Result { + self.check_mutable().await?; + let schema = self.schema().await?; + let plan = crate::table::computed_columns::plan_function_application( + schema.as_ref(), + application, + output_name, + )?; + let new_columns = plan + .outputs + .iter() + .map(|output| { + serde_json::json!({ + "name": output.output_name, + "all_null": true, + }) + }) + .collect::>(); + let mut body = serde_json::json!({ + "new_columns": new_columns, + "function": { + "application": plan.application, + "binding_metadata_version": plan.binding_metadata_version, + "input_bindings": plan.input_bindings, + "input_schema": plan.input_schema, + "output_schema": plan.output_schema, + "outputs": plan.outputs, + }, + }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/add_columns/", self.identifier)) + .json(&body); + let freshness_request = self.snapshot_freshness_headers(); + let (request_id, response) = self + .send_with_freshness(request, true, freshness_request) + .await?; + // A Function declaration can return 404 for the Function rather than the table. + let response = self + .client + .check_response(&request_id, response) + .await + .inspect_err(|error| self.handle_error_invalidation(error))?; + let body = response.text().await.err_to_http(request_id.clone())?; + + if body.trim().is_empty() { + return Ok(AddColumnsResult { version: 0 }); + } + + let result: AddColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse add Function columns response: {e}").into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(freshness_request, result.version); + Ok(result) + } + + pub(super) async fn refresh_column_async_impl( + &self, + column: &str, + ) -> Result> { + self.check_mutable().await?; + let mut body = serde_json::json!({ "column": column }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/backfill_column", self.identifier)) + .json(&body); + let (request_id, response) = self.send(request, true).await?; + // Preserve dependency errors: a deleted bound Function also returns 404. + let response = self + .client + .check_response(&request_id, response) + .await + .inspect_err(|error| self.handle_error_invalidation(error))?; + let body = response.text().await.err_to_http(request_id.clone())?; + + #[derive(serde::Deserialize)] + struct BackfillResponse { + job_id: String, + } + let response: BackfillResponse = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse backfill_column response: {}", e).into(), + request_id, + status_code: None, + })?; + + Ok(Job::new_typed(Box::new(FreshnessJob { + inner: RemoteJob::new(self.client.clone(), response.job_id), + freshness: self.freshness.clone(), + version: self.version.clone(), + tracked_result: TrackedJobResult::RefreshColumn, + freshness_request: self.snapshot_freshness_headers(), + }))) + } + + pub(super) async fn function_errors_impl( + &self, + request: &crate::function::FunctionErrorsRequest, + ) -> Result { + let mut body = serde_json::json!({}); + if let Some(job_id) = &request.job_id { + body["job_id"] = serde_json::json!(job_id); + } + if let Some(column) = &request.column { + body["column"] = serde_json::json!(column); + } + if let Some(limit) = request.limit { + body["limit"] = serde_json::json!(limit); + } + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/errors", self.identifier)) + .json(&body); + let (request_id, response) = self.send(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse errors response: {}", e).into(), + request_id, + status_code: None, + }) + } + + pub(super) async fn alter_columns_impl( + &self, + alterations: &[ColumnAlteration], + ) -> Result { + self.check_mutable().await?; + let body = alterations + .iter() + .map(|alteration| { + let mut value = serde_json::json!({ + "path": alteration.path, + }); + if let Some(rename) = &alteration.rename { + value["rename"] = serde_json::Value::String(rename.clone()); + } + if let Some(data_type) = &alteration.data_type { + let json_data_type = JsonDataType::try_from(data_type).unwrap(); + let json_data_type = serde_json::to_value(&json_data_type).unwrap(); + value["data_type"] = json_data_type; + } + if let Some(nullable) = &alteration.nullable { + value["nullable"] = serde_json::Value::Bool(*nullable); + } + value + }) + .collect::>(); + let mut body = serde_json::json!({ "alterations": body }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/alter_columns/", self.identifier)) + .json(&body); + let freshness_request = self.snapshot_freshness_headers(); + let (request_id, response) = self + .send_with_freshness(request, true, freshness_request) + .await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + if body.trim().is_empty() { + // Backward compatible with old servers + return Ok(AlterColumnsResult { version: 0 }); + } + + let result: AlterColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse alter_columns response: {}", e).into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(freshness_request, result.version); + + Ok(result) + } + + pub(super) async fn update_field_metadata_impl( + &self, + updates: &[FieldMetadataUpdate], + ) -> Result { + self.check_mutable().await?; + let mut body = serde_json::json!({ "updates": updates }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!( + "/v1/table/{}/update_field_metadata/", + self.identifier + )) + .json(&body); + let freshness_request = self.snapshot_freshness_headers(); + let (request_id, response) = self + .send_with_freshness(request, true, freshness_request) + .await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + let result: UpdateFieldMetadataResult = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse update_field_metadata response: {}", e).into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(freshness_request, result.version); + Ok(result) + } + + pub(super) async fn drop_columns_impl(&self, columns: &[&str]) -> Result { + self.check_mutable().await?; + let mut body = serde_json::json!({ "columns": columns }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!("/v1/table/{}/drop_columns/", self.identifier)) + .json(&body); + let freshness_request = self.snapshot_freshness_headers(); + let (request_id, response) = self + .send_with_freshness(request, true, freshness_request) + .await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + if body.trim().is_empty() { + // Backward compatible with old servers + return Ok(DropColumnsResult { version: 0 }); + } + + let result: DropColumnsResult = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse drop_columns response: {}", e).into(), + request_id, + status_code: None, + })?; + + self.invalidate_schema_cache(); + self.track_write_version(freshness_request, result.version); + + Ok(result) + } +} diff --git a/rust/lancedb/src/remote/table/describe.rs b/rust/lancedb/src/remote/table/describe.rs new file mode 100644 index 000000000..75a1ca615 --- /dev/null +++ b/rust/lancedb/src/remote/table/describe.rs @@ -0,0 +1,142 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[derive(Deserialize)] +pub(super) struct TableDescription { + pub(super) version: u64, + pub(super) schema: JsonSchema, + pub(super) location: Option, +} + +/// How a response body frames its Arrow IPC payload. `/query` answers with file framing +/// and `fetch_blobs` with stream framing, so the reader is picked per response. +pub(super) enum ArrowIpcFraming { + File, + Stream, +} + +/// A response with no `Content-Type` uses file framing, preserving this helper's +/// behavior before `fetch_blobs` introduced stream responses. +pub(super) fn resolve_arrow_ipc_framing( + content_type: Option<&str>, + request_id: &str, +) -> Result { + let Some(media_type) = content_type.map(base_media_type) else { + return Ok(ArrowIpcFraming::File); + }; + if media_type.eq_ignore_ascii_case(ARROW_STREAM_CONTENT_TYPE) { + return Ok(ArrowIpcFraming::Stream); + } + if media_type.eq_ignore_ascii_case(ARROW_FILE_CONTENT_TYPE) { + return Ok(ArrowIpcFraming::File); + } + Err(Error::Http { + source: format!( + "Expected an Arrow IPC response with Content-Type '{ARROW_STREAM_CONTENT_TYPE}' \ + or '{ARROW_FILE_CONTENT_TYPE}', got '{media_type}'" + ) + .into(), + request_id: request_id.into(), + status_code: None, + }) +} + +/// An Arrow IPC stream carrying `schema` and no batches. +/// +/// The body of an all-null `add_columns`: the server reads the fields to +/// add straight off the schema message, so nothing but the field list is +/// worth sending. +pub(super) fn write_ipc_schema(schema: &arrow_schema::Schema) -> Result> { + let mut body = Vec::new(); + { + let mut writer = arrow_ipc::writer::StreamWriter::try_new(&mut body, schema)?; + writer.finish()?; + } + Ok(body) +} + +/// Strip media-type parameters before matching against the Arrow content types. +pub(super) fn base_media_type(content_type: &str) -> &str { + match content_type.split_once(';') { + Some((media_type, _parameters)) => media_type.trim(), + None => content_type.trim(), + } +} + +/// Extract an Error from Arc, reconstructing if the Arc is shared. +/// This is needed because `Shared` futures cache results internally, so +/// `Arc::try_unwrap` typically fails. +pub(super) fn unwrap_shared_error(arc: Arc) -> Error { + match Arc::try_unwrap(arc) { + Ok(err) => err, + Err(arc) => match &*arc { + Error::TableNotFound { name, source } => Error::TableNotFound { + name: name.clone(), + source: source.to_string().into(), + }, + _ => Error::Runtime { + message: arc.to_string(), + }, + }, + } +} + +pub(super) async fn fetch_schema( + client: &RestfulLanceDbClient, + identifier: &str, + table_name: &str, + read_snapshot: ReadSnapshot, + branch: Option, + freshness: Arc>, +) -> Result { + let mut body = serde_json::json!({ "version": read_snapshot.version }); + if let Some(branch) = &branch { + body["branch"] = serde_json::Value::String(branch.clone()); + } + let freshness_headers = read_snapshot.freshness; + let request = freshness_headers + .apply(client.post(&format!("/v1/table/{}/describe/", identifier))) + .json(&body); + + let (request_id, response) = client.send_with_retry(request, None, true).await?; + + if response.status() == StatusCode::NOT_FOUND { + let body = response.text().await.ok().unwrap_or_default(); + return Err(Error::TableNotFound { + name: table_name.to_string(), + source: Box::new(Error::Http { + source: body.into(), + request_id, + status_code: Some(StatusCode::NOT_FOUND), + }), + }); + } + + let response = client.check_response(&request_id, response).await?; + freshness_headers.observe_headers(&freshness, response.headers()); + let body = response.text().await.map_err(|e| { + let status_code = e.status(); + Error::Http { + source: Box::new(e), + request_id: request_id.clone(), + status_code, + } + })?; + + let description: TableDescription = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse table description: {}", e).into(), + request_id, + status_code: None, + })?; + freshness_headers.observe_version(&freshness, description.version); + if !freshness_headers.is_current(&freshness) { + return Err(Error::Runtime { + message: SCHEMA_SELECTOR_CHANGED.to_string(), + }); + } + + let arrow_schema: arrow_schema::Schema = description.schema.try_into()?; + Ok(Arc::new(arrow_schema)) +} diff --git a/rust/lancedb/src/remote/table/freshness.rs b/rust/lancedb/src/remote/table/freshness.rs new file mode 100644 index 000000000..90b2eb302 --- /dev/null +++ b/rust/lancedb/src/remote/table/freshness.rs @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +/// Per-table state driving the freshness headers (`x-lancedb-min-version`, +/// `x-lancedb-min-timestamp`, and `x-lancedb-min-read-version`) sent on table +/// requests. +#[derive(Debug, Default, Clone, Copy)] +pub(super) struct FreshnessState { + /// Identifies the handle timeline that produced this state. Explicit + /// checkout operations advance the generation so responses from older + /// in-flight requests cannot repopulate the new timeline's constraints. + pub(super) generation: u64, + /// Exact-version, tag, and snapshot handles must not carry latest-timeline + /// constraints. Their request body already selects the precise version. + pub(super) pinned: bool, + /// Provides read-your-write within a single handle: writes that return a + /// version update this, and reads send it as `x-lancedb-min-version`. + pub(super) min_version: Option, + /// Highest committed dataset version advertised by a successful table + /// response on this handle. Later requests send it as + /// `x-lancedb-min-read-version` so a load-balanced query + /// node whose cache is behind this version must refresh before serving, + /// giving monotonic reads across nodes regardless of which one the load + /// balancer routes to. Unlike write result bodies, this is sourced only + /// from the server's committed dataset-version response header or other + /// typed dataset-version fields, so WAL entry ids cannot enter it. + pub(super) min_read_version: Option, + /// Wall-clock time captured at the last [`BaseTable::checkout_latest`] + /// call. Subsequent reads send + /// `max(baseline, now - read_consistency_interval)` as + /// `x-lancedb-min-timestamp`. + /// + /// Without this, `checkout_latest()` would have no effect on subsequent + /// reads when `read_consistency_interval` is unset (the default): a + /// server-side cache could still serve a snapshot older than the moment + /// the user explicitly asked for "latest". The baseline forces the + /// server to skip any cache entry older than the checkout time, so the + /// `checkout_latest()` signal is preserved across reads on the same + /// handle regardless of the configured consistency interval. + pub(super) checkout_baseline: Option, +} + +/// Snapshot of the headers that should be attached to a single table request. +#[derive(Debug, Default, Clone, Copy)] +pub(super) struct FreshnessHeaders { + generation: u64, + min_version: Option, + pub(super) min_timestamp: Option, + min_read_version: Option, +} + +#[derive(Debug, Default, Clone, Copy)] +pub(super) struct ReadSnapshot { + pub(super) version: Option, + pub(super) freshness_state: FreshnessState, + pub(super) freshness: FreshnessHeaders, +} + +impl FreshnessHeaders { + pub(super) fn apply(self, mut request: RequestBuilder) -> RequestBuilder { + if let Some(v) = self.min_version { + request = request.header(MIN_VERSION_HEADER, v.to_string()); + } + if let Some(ts) = self.min_timestamp { + let dt: chrono::DateTime = ts.into(); + request = request.header(MIN_TIMESTAMP_HEADER, dt.to_rfc3339()); + } + if let Some(v) = self.min_read_version { + request = request.header(MIN_READ_VERSION_HEADER, v.to_string()); + } + request + } + + pub(super) fn observe_version(self, freshness: &Mutex, version: u64) { + track_read_version_for_generation(freshness, self.generation, version); + } + + pub(super) fn observe_headers( + self, + freshness: &Mutex, + headers: &reqwest::header::HeaderMap, + ) { + if let Some(version) = headers + .get(&VERSION_HEADER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + { + self.observe_version(freshness, version); + } + } + + pub(super) fn update_if_current( + self, + freshness: &Mutex, + update: impl FnOnce(&mut FreshnessState), + ) { + let mut state = freshness.lock().unwrap(); + if state.generation == self.generation { + update(&mut state); + } + } + + pub(super) fn is_current(self, freshness: &Mutex) -> bool { + freshness.lock().unwrap().generation == self.generation + } +} + +pub(super) fn track_read_version(freshness: &Mutex, version: u64) { + if version == 0 { + return; + } + let mut state = freshness.lock().unwrap(); + state.min_read_version = Some(state.min_read_version.map_or(version, |v| v.max(version))); +} + +pub(super) fn track_read_version_for_generation( + freshness: &Mutex, + generation: u64, + version: u64, +) { + if version == 0 { + return; + } + let mut state = freshness.lock().unwrap(); + if state.generation == generation { + state.min_read_version = Some(state.min_read_version.map_or(version, |v| v.max(version))); + } +} + +/// A backfill job whose successful wait establishes a read-freshness +/// baseline on the submitting handle, so a later read cannot be served +/// from a cache older than the completed fill. A handle pinned by checkout +/// at completion keeps its time-travel view instead. +pub(super) struct FreshnessJob { + pub(super) inner: RemoteJob, + pub(super) freshness: Arc>, + pub(super) version: Arc>>, + pub(super) tracked_result: TrackedJobResult, + pub(super) freshness_request: FreshnessHeaders, +} + +#[derive(Clone, Copy)] +pub(super) enum TrackedJobResult { + None, + RefreshColumn, +} + +#[async_trait] +impl crate::job::JobHandle for FreshnessJob { + fn id(&self) -> Option<&str> { + crate::job::JobHandle::id(&self.inner) + } + + async fn status(&self) -> Result { + crate::job::JobHandle::status(&self.inner).await + } + + async fn describe(&self) -> Result { + crate::job::JobHandle::describe(&self.inner).await + } + + async fn events( + &self, + request: crate::job::JobEventsRequest, + ) -> Result> { + crate::job::JobHandle::events(&self.inner, request).await + } + + async fn wait(&self) -> Result { + let result = crate::job::JobHandle::wait(&self.inner).await?; + let version = self.version.read().await; + if version.is_none() { + let result_version = match self.tracked_result { + TrackedJobResult::None => None, + TrackedJobResult::RefreshColumn => result.value().and_then(|value| { + serde_json::from_value::(value.clone()) + .ok() + .map(|result| { + result + .published_version + .map_or(result.source_version, |version| { + version.max(result.source_version) + }) + }) + }), + } + .filter(|version| *version != 0); + if let Some(version) = result_version { + self.freshness_request + .observe_version(&self.freshness, version); + } else { + self.freshness_request + .update_if_current(&self.freshness, |state| { + state.checkout_baseline = Some(SystemTime::now()); + }); + } + } + Ok(result) + } + + async fn cancel(&self) -> Result<()> { + crate::job::JobHandle::cancel(&self.inner).await + } +} + +pub(super) fn compute_min_timestamp( + state: &FreshnessState, + interval: Option, + now: SystemTime, +) -> Option { + let interval_based = match interval { + None => None, + Some(d) if d.is_zero() => Some(now), + Some(d) => Some(now.checked_sub(d).unwrap_or(now)), + }; + match (interval_based, state.checkout_baseline) { + (None, None) => None, + (Some(t), None) | (None, Some(t)) => Some(t), + (Some(a), Some(b)) => Some(a.max(b)), + } +} + +pub(super) fn freshness_headers_snapshot( + freshness: &Mutex, + interval: Option, +) -> FreshnessHeaders { + freshness_state_snapshot(freshness, interval).1 +} + +pub(super) fn freshness_state_snapshot( + freshness: &Mutex, + interval: Option, +) -> (FreshnessState, FreshnessHeaders) { + let state = *freshness.lock().unwrap(); + if state.pinned { + return ( + state, + FreshnessHeaders { + generation: state.generation, + ..FreshnessHeaders::default() + }, + ); + } + ( + state, + FreshnessHeaders { + generation: state.generation, + min_version: state.min_version, + min_timestamp: compute_min_timestamp(&state, interval, SystemTime::now()), + min_read_version: state.min_read_version, + }, + ) +} diff --git a/rust/lancedb/src/remote/table/indices.rs b/rust/lancedb/src/remote/table/indices.rs new file mode 100644 index 000000000..7f45b17a1 --- /dev/null +++ b/rust/lancedb/src/remote/table/indices.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +/// Deserialize an index's `created_at` field. +/// +/// The server returns this as an RFC 3339 string (e.g. `"2026-06-18T21:37:36.637Z"`), +/// but older deployments sent a unix timestamp in milliseconds. Accept both so the +/// client works against any server version. +pub(super) fn deserialize_created_at<'de, D>( + deserializer: D, +) -> std::result::Result>, D::Error> +where + D: serde::Deserializer<'de>, +{ + use serde::de::Error as _; + + #[derive(Deserialize)] + #[serde(untagged)] + enum CreatedAt { + Rfc3339(String), + Millis(i64), + } + + match Option::::deserialize(deserializer)? { + None => Ok(None), + Some(CreatedAt::Rfc3339(s)) => DateTime::parse_from_rfc3339(&s) + .map(|dt| Some(dt.with_timezone(&Utc))) + .map_err(D::Error::custom), + Some(CreatedAt::Millis(ms)) => Ok(DateTime::from_timestamp_millis(ms)), + } +} + +impl RemoteTable { + pub(super) async fn index_stats_read_snapshot( + &self, + index_name: &str, + read_snapshot: ReadSnapshot, + ) -> Result> { + let encoded_name = urlencoding::encode(index_name); + let mut body = serde_json::json!({ "version": read_snapshot.version }); + self.apply_branch_body(&mut body); + let request = self + .client + .post(&format!( + "/v1/table/{}/index/{encoded_name}/stats/", + self.identifier + )) + .json(&body); + + let (request_id, response) = self + .send_with_freshness(request, true, read_snapshot.freshness) + .await?; + if response.status() == StatusCode::NOT_FOUND { + return Ok(None); + } + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let stats = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse index statistics: {}", e).into(), + request_id, + status_code: None, + })?; + Ok(Some(stats)) + } + + /// Parse the response from `/index/list/` into `IndexConfig` entries. + /// + /// When the server returns `index_type` inline, all enriched fields are + /// used directly and no further requests are made. When `index_type` is + /// absent (legacy servers), a `/index/{name}/stats/` call is made for each + /// index to retrieve the type. + pub(super) async fn parse_index_list_response( + &self, + body: &str, + request_id: &str, + schema: &SchemaRef, + read_snapshot: ReadSnapshot, + ) -> Result> { + use crate::index::IndexType; + + #[derive(Deserialize)] + struct ListIndicesResponse { + indexes: Vec, + } + + #[derive(Deserialize)] + struct IndexListEntry { + index_name: String, + columns: Vec, + // Present on enriched responses; absent on legacy servers. + // Used as the sentinel to decide whether to skip the stats call. + index_type: Option, + index_uuid: Option, + #[serde(default, deserialize_with = "deserialize_created_at")] + created_at: Option>, + num_indexed_rows: Option, + num_unindexed_rows: Option, + size_bytes: Option, + num_segments: Option, + index_version: Option, + index_details: Option, + type_url: Option, + } + + let response: ListIndicesResponse = + serde_json::from_str(body).map_err(|err| Error::Http { + source: format!( + "Failed to parse list_indices response: {}, body: {}", + err, body + ) + .into(), + request_id: request_id.to_string(), + status_code: None, + })?; + + let mut futures = Vec::with_capacity(response.indexes.len()); + for entry in response.indexes { + let columns = entry + .columns + .iter() + .map(|column| { + resolve_arrow_field_path(schema, column) + .map(|(canonical_column, _)| canonical_column) + }) + .collect::>>()?; + + let future = async move { + if let Some(index_type) = entry.index_type { + // Enriched response: all fields available, no stats call needed. + Ok(Some(IndexConfig { + name: entry.index_name, + index_type, + columns, + index_uuid: entry.index_uuid, + type_url: entry.type_url, + created_at: entry.created_at, + num_indexed_rows: entry.num_indexed_rows, + num_unindexed_rows: entry.num_unindexed_rows, + size_bytes: entry.size_bytes, + num_segments: entry.num_segments, + index_version: entry.index_version, + index_details: entry.index_details, + })) + } else { + // Legacy response: fetch index type via stats endpoint. + match self + .index_stats_read_snapshot(&entry.index_name, read_snapshot) + .await + { + Ok(Some(stats)) => Ok(Some(IndexConfig { + name: entry.index_name, + index_type: stats.index_type, + columns, + index_uuid: None, + type_url: None, + created_at: None, + num_indexed_rows: None, + num_unindexed_rows: None, + size_bytes: None, + num_segments: None, + index_version: None, + index_details: None, + })), + Ok(None) => Ok(None), // Index deleted since we listed it. + Err(e) => Err(e), + } + } + }; + futures.push(future); + } + + let results = futures::future::try_join_all(futures).await?; + let mut indices: Vec = results.into_iter().flatten().collect(); + let lance_schema = lance_core::datatypes::Schema::try_from(schema.as_ref())?; + for index in &mut indices { + if index.index_type == IndexType::FTS { + // The wire format uses physical paths for schema resolution. Match + // native tables by exposing list-transparent paths to callers. + for column in &mut index.columns { + let field_id = lance_schema.field_id(column)?; + *column = public_fts_field_path_by_id(&lance_schema, field_id)?; + } + } + } + Ok(indices) + } +} diff --git a/rust/lancedb/src/remote/table/lsm.rs b/rust/lancedb/src/remote/table/lsm.rs new file mode 100644 index 000000000..04c6c3b7a --- /dev/null +++ b/rust/lancedb/src/remote/table/lsm.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +impl RemoteTable { + pub(super) async fn get_lsm_stats_impl( + &self, + include_generation_rows: bool, + ) -> Result> { + // Read-semantics POST, like `get_lsm_write_spec`. + let request = self + .client + .post(&format!("/v1/table/{}/get_lsm_stats/", self.identifier)) + .json(&serde_json::json!({ + "include_generation_rows": include_generation_rows, + })); + let (request_id, response) = self.send_lsm_route(request).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let parsed: GetLsmStatsResponse = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse get_lsm_stats response: {e}").into(), + request_id, + status_code: None, + })?; + // `null` — and only — when the table has no LSM write path. + Ok(parsed.lsm_stats) + } + + pub(super) async fn set_lsm_write_spec_impl(&self, spec: LsmWriteSpec) -> Result<()> { + self.check_mutable().await?; + + // Map the spec onto the server's request DTO. `sharding` is internally + // tagged on `mode` to mirror sophon's `Sharding` enum. A null + // `maintained_indexes` asks the server to resolve every maintainable + // index at HEAD; a list is verbatim, an empty one meaning none. + let sharding = match &spec { + LsmWriteSpec::Bucket { + column, + num_buckets, + .. + } => serde_json::json!({ + "mode": "bucket", + "column": column, + "num_buckets": num_buckets, + }), + LsmWriteSpec::Identity { column, .. } => serde_json::json!({ + "mode": "identity", + "column": column, + }), + LsmWriteSpec::Unsharded { .. } => serde_json::json!({ "mode": "unsharded" }), + }; + let body = serde_json::json!({ + "sharding": sharding, + "maintained_indexes": spec.maintained_indexes(), + "writer_config_defaults": spec.writer_config_defaults(), + }); + + let request = self + .client + .post(&format!( + "/v1/table/{}/set_lsm_write_spec/", + self.identifier + )) + .json(&body); + let (request_id, response) = self.send(request, true).await?; + self.check_table_response(&request_id, response).await?; + Ok(()) + } + + pub(super) async fn unset_lsm_write_spec_impl(&self) -> Result<()> { + self.check_mutable().await?; + let request = self.client.post(&format!( + "/v1/table/{}/unset_lsm_write_spec/", + self.identifier + )); + let (request_id, response) = self.send(request, true).await?; + self.check_table_response(&request_id, response).await?; + self.wal_pk_fusion.forget(); // WAL-PK-FUSION: delete. + Ok(()) + } + + pub(super) async fn get_lsm_write_spec_impl(&self) -> Result> { + // Read counterpart to set/unset, resolved server-side against HEAD. The + // server reads the spec from the `__lance_mem_wal` system index (shard + // column mapped from its Lance field id against the current schema) and + // re-encodes it into the same sophon-owned shape the set endpoint + // accepts — no lance/lancedb types cross the wire. `lsm_write_spec` is + // null when the LSM write path is not enabled for the table. + let request = self.client.post(&format!( + "/v1/table/{}/get_lsm_write_spec/", + self.identifier + )); + let (request_id, response) = self.send(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + + // Mirror of sophon's `Sharding` (internally tagged on `mode`) and + // `LsmWriteSpecBody` / `GetLsmWriteSpecResponse`. + #[derive(Deserialize)] + #[serde(tag = "mode", rename_all = "snake_case")] + enum Sharding { + Unsharded, + Bucket { column: String, num_buckets: u32 }, + Identity { column: String }, + } + #[derive(Deserialize)] + struct LsmWriteSpecBody { + sharding: Sharding, + /// `null` selects every index the table has; `[]` selects none. + #[serde(default)] + maintained_indexes: Option>, + #[serde(default)] + writer_config_defaults: std::collections::HashMap, + } + #[derive(Deserialize)] + struct GetLsmWriteSpecResponse { + lsm_write_spec: Option, + } + + let parsed: GetLsmWriteSpecResponse = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse get_lsm_write_spec response: {}", e).into(), + request_id, + status_code: None, + })?; + + let Some(body) = parsed.lsm_write_spec else { + // The LSM write path is not enabled for this table. + return Ok(None); + }; + + let spec = match body.sharding { + Sharding::Bucket { + column, + num_buckets, + } => LsmWriteSpec::bucket(column, num_buckets), + Sharding::Identity { column } => LsmWriteSpec::identity(column), + Sharding::Unsharded => LsmWriteSpec::unsharded(), + } + .with_maintained_indexes(body.maintained_indexes) + .with_writer_config_defaults(body.writer_config_defaults); + + Ok(Some(spec)) + } +} diff --git a/rust/lancedb/src/remote/table/merge.rs b/rust/lancedb/src/remote/table/merge.rs new file mode 100644 index 000000000..afcd7fc59 --- /dev/null +++ b/rust/lancedb/src/remote/table/merge.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +#[derive(Serialize, Clone, Debug)] +pub struct MergeInsertRequest { + // Sent as one repeated `on` query parameter per column, which is how the + // namespace spec encodes an array-valued `on`. serde_urlencoded (which + // reqwest's `query()` uses) cannot serialize a sequence nested in a struct, + // so this field is emitted separately by [`Self::on_query_params`]. + #[serde(skip_serializing)] + pub(super) on: Vec, + pub(super) when_matched_update_all: bool, + pub(super) when_matched_update_all_filt: Option, + pub(super) when_not_matched_insert_all: bool, + pub(super) when_not_matched_by_source_delete: bool, + pub(super) when_not_matched_by_source_delete_filt: Option, + // For backwards compatibility, only serialize use_index when it's false + // (the default is true) + #[serde(skip_serializing_if = "is_true")] + pub(super) use_index: bool, + // Only serialize use_lsm when explicitly set (Some); a server that predates + // it ignores the field and routes as it would by default. + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) use_lsm: Option, +} + +impl MergeInsertRequest { + /// The `on` columns as repeated query parameters: `?on=a&on=b`. + /// + /// A single column serializes to `?on=a`, exactly what clients sent before + /// `on` became a list, so a server that predates composite keys sees no + /// change from a single-column caller. + pub(crate) fn on_query_params(&self) -> Vec<(&str, &str)> { + self.on.iter().map(|col| ("on", col.as_str())).collect() + } +} + +pub(super) fn is_true(b: &bool) -> bool { + *b +} + +impl TryFrom for MergeInsertRequest { + type Error = Error; + + fn try_from(value: MergeInsertBuilder) -> Result { + if value.on.is_empty() { + return Err(Error::InvalidInput { + message: "MergeInsertBuilder missing required 'on' field".into(), + }); + } + // The server rejects a repeated column with a 400; catching it here + // names the offending column and costs no round trip. + let mut seen = HashSet::with_capacity(value.on.len()); + if let Some(dup) = value.on.iter().find(|col| !seen.insert(*col)) { + return Err(Error::InvalidInput { + message: format!("MergeInsertBuilder 'on' column '{dup}' is repeated"), + }); + } + + let when_matched_update_all_filt = match value.when_matched_update_all_filt { + Some(MergeFilter::Sql(sql)) => Some(sql), + Some(MergeFilter::Expr(_)) => { + return Err(Error::NotSupported { + message: "DataFusion expressions are not supported on remote tables".into(), + }); + } + None => None, + }; + + let when_not_matched_by_source_delete_filt = + match value.when_not_matched_by_source_delete_filt { + Some(MergeFilter::Sql(sql)) => Some(sql), + Some(MergeFilter::Expr(_)) => { + return Err(Error::NotSupported { + message: "DataFusion expressions are not supported on remote tables".into(), + }); + } + None => None, + }; + + Ok(Self { + on: value.on, + when_matched_update_all: value.when_matched_update_all, + when_matched_update_all_filt, + when_not_matched_insert_all: value.when_not_matched_insert_all, + when_not_matched_by_source_delete: value.when_not_matched_by_source_delete, + when_not_matched_by_source_delete_filt, + // Only serialize use_index when it's false for backwards compatibility + use_index: value.use_index, + use_lsm: value.use_lsm, + }) + } +} diff --git a/rust/lancedb/src/remote/table/requests.rs b/rust/lancedb/src/remote/table/requests.rs new file mode 100644 index 000000000..a67c011b2 --- /dev/null +++ b/rust/lancedb/src/remote/table/requests.rs @@ -0,0 +1,958 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +impl RemoteTable { + pub(super) async fn submit_create_index( + &self, + mut index: IndexBuilder, + ) -> Result> { + self.check_mutable().await?; + let request = self + .client + .post(&format!("/v1/table/{}/create_index/", self.identifier)); + + let column = match index.columns.len() { + 0 => { + return Err(Error::InvalidInput { + message: "No columns specified".into(), + }); + } + 1 => index.columns.pop().unwrap(), + _ => { + return Err(Error::NotSupported { + message: "Indices over multiple columns not yet supported".into(), + }); + } + }; + if matches!( + &index.index, + Index::FTS(params) if params.get_document_granularity().is_list_element() + ) && !self.server_version.support_fts_document_granularity() + { + return Err(Error::NotSupported { + message: "FTS document granularity requires remote server version 0.6.0 or later" + .into(), + }); + } + let schema = self.schema().await?; + let (canonical_column, field) = match &index.index { + Index::FTS(_) => resolve_arrow_fts_field_path(&schema, &column)?, + _ => resolve_arrow_field_path(&schema, &column)?, + }; + let mut body = serde_json::json!({ + "column": canonical_column + }); + + if !index.replace { + body["replace"] = false.into(); + } + + // Add name parameter if provided (for backwards compatibility, only include if Some) + if let Some(ref name) = index.name { + body["name"] = serde_json::Value::String(name.clone()); + } + + // Warn if train=false is specified since it's not meaningful + if !index.train { + log::warn!( + "train=false has no effect remote tables. The index will be created empty and automatically populated in the background." + ); + } + + fn to_json(params: &impl serde::Serialize) -> crate::Result { + serde_json::to_value(params).map_err(|e| Error::InvalidInput { + message: format!("failed to serialize index params {:?}", e), + }) + } + + // Map each Index variant to its wire type name and serializable params. + // Auto is special-cased since it needs schema inspection. + let (index_type_str, params) = match &index.index { + Index::IvfFlat(p) => ("IVF_FLAT", Some(to_json(p)?)), + Index::IvfPq(p) => ("IVF_PQ", Some(to_json(p)?)), + Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)), + Index::IvfHnswSq(p) => ("IVF_HNSW_SQ", Some(to_json(p)?)), + Index::IvfHnswFlat(p) => ("IVF_HNSW_FLAT", Some(to_json(p)?)), + Index::IvfRq(p) => ("IVF_RQ", Some(to_json(p)?)), + Index::BTree(p) => ("BTREE", Some(to_json(p)?)), + Index::Bitmap(p) => ("BITMAP", Some(to_json(p)?)), + Index::LabelList(p) => ("LABEL_LIST", Some(to_json(p)?)), + Index::Fm(p) => ("FM", Some(to_json(p)?)), + Index::ZoneMap(p) => ("ZONEMAP", Some(to_json(p)?)), + Index::NGram(p) => ("NGRAM", Some(to_json(p)?)), + Index::BloomFilter(p) => ("BLOOM_FILTER", Some(to_json(p)?)), + Index::RTree(p) => ("RTREE", Some(to_json(p)?)), + Index::FTS(p) => { + let mut params = to_json(p)?; + if p.get_document_granularity().is_list_element() { + params["document_granularity"] = "list_element".into(); + } + ("FTS", Some(params)) + } + Index::Auto => { + if supported_vector_data_type(field.data_type()) { + body[METRIC_TYPE_KEY] = + serde_json::Value::String(DistanceType::L2.to_string().to_lowercase()); + ("IVF_PQ", None) + } else if supported_btree_data_type(field.data_type()) { + ("BTREE", None) + } else { + return Err(Error::NotSupported { + message: format!( + "there are no indices supported for the field `{}` with the data type {}", + field.name(), + field.data_type() + ), + }); + } + } + _ => { + return Err(Error::NotSupported { + message: "Index type not supported".into(), + }); + } + }; + + body[INDEX_TYPE_KEY] = index_type_str.into(); + if let Some(params) = params { + for (key, value) in params.as_object().expect("params should be a JSON object") { + body[key] = value.clone(); + } + } + self.apply_branch_body(&mut body); + + let request = request.json(&body); + + let (request_id, response) = self.send(request, true).await?; + + let response = self.check_table_response(&request_id, response).await?; + let job_id = response + .text() + .await + .ok() + .and_then(|body| extract_job_id(&body)); + + if let Some(wait_timeout) = index.wait_timeout { + let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); + self.wait_for_index(&[&index_name], wait_timeout).await?; + } + + Ok(job_id) + } + + pub(in crate::remote) fn new_with_sql_client( + client: RestfulLanceDbClient, + name: String, + namespace: Vec, + identifier: String, + server_version: ServerVersion, + sql_client: Option, + ) -> Self { + Self { + client, + name, + namespace, + identifier, + server_version, + sql_client, + version: Arc::new(RwLock::new(None)), + location: RwLock::new(None), + schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + wal_pk_fusion: PkFusionMemory::default(), + freshness: Arc::new(Mutex::new(FreshnessState::default())), + branch: None, + } + } + + /// Seed the schema cache from a `describe` body the caller already fetched. + /// + /// Best effort. `open_table` succeeds today without reading this body, so a + /// body we cannot parse leaves the cache empty and the next schema read + /// fetches it again through the path that reports a real error. + pub(crate) fn seed_schema(&self, describe_body: &str) { + let Ok(description) = serde_json::from_str::(describe_body) else { + return; + }; + self.track_read_version(description.version); + if let Ok(schema) = arrow_schema::Schema::try_from(description.schema) { + self.schema_cache.seed(Arc::new(schema)); + } + } + + /// Return a new handle scoped to `branch`, sharing the client but with fresh + /// caches and version/freshness state (the branch tracks its own latest). + /// Mirrors `NativeTable`'s handle-per-branch model. + pub(super) fn with_branch(&self, branch: Option) -> Self { + Self { + client: self.client.clone(), + name: self.name.clone(), + namespace: self.namespace.clone(), + identifier: self.identifier.clone(), + server_version: self.server_version.clone(), + sql_client: self.sql_client.clone(), + version: Arc::new(RwLock::new(None)), + location: RwLock::new(None), + schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + wal_pk_fusion: PkFusionMemory::default(), + freshness: Arc::new(Mutex::new(FreshnessState::default())), + branch, + } + } + + /// Stamp the branch onto a request as a `?branch=` query param (used for + /// Arrow-body / query-only ops). `None` (main) leaves the request unchanged, + /// keeping it byte-identical to the non-branch path. + pub(super) fn apply_branch_query(&self, request: RequestBuilder) -> RequestBuilder { + match &self.branch { + Some(branch) => request.query(&[("branch", branch.as_str())]), + None => request, + } + } + + /// Stamp the branch into a JSON request body under `"branch"` (used for JSON + /// ops). `None` (main) leaves the body unchanged. + pub(super) fn apply_branch_body(&self, body: &mut serde_json::Value) { + if let Some(branch) = &self.branch { + body["branch"] = serde_json::Value::String(branch.clone()); + } + } + + pub(super) async fn describe(&self) -> Result { + self.describe_read_snapshot(self.snapshot_read_state().await) + .await + } + + pub(super) async fn describe_read_snapshot( + &self, + read_snapshot: ReadSnapshot, + ) -> Result { + let request = self + .client + .post(&format!("/v1/table/{}/describe/", self.identifier)); + self.describe_with_request( + request, + read_snapshot.version, + Some(read_snapshot.freshness), + ) + .await + } + + pub(super) async fn schema_read_snapshot( + &self, + read_snapshot: ReadSnapshot, + ) -> Result { + if read_snapshot.freshness.is_current(&self.freshness) + && let Some(schema) = self.schema_cache.try_get() + && read_snapshot.freshness.is_current(&self.freshness) + { + return Ok(schema); + } + + let description = self.describe_read_snapshot(read_snapshot).await?; + Ok(Arc::new(description.schema.try_into()?)) + } + + pub(super) async fn resolve_tag_version_with_request( + &self, + tag: &str, + request: RequestBuilder, + fenced: bool, + ) -> Result { + let request = request.json(&serde_json::json!({ "tag": tag })); + + let (request_id, response) = if fenced { + self.send(request, true).await? + } else { + self.send_unfenced(request, true).await? + }; + let response = self.check_table_response(&request_id, response).await?; + + match response.text().await { + Ok(body) => { + let value: serde_json::Value = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse tag version: {}", e).into(), + request_id: request_id.clone(), + status_code: None, + })?; + + value + .get("version") + .and_then(|v| v.as_u64()) + .ok_or_else(|| Error::Http { + source: format!("Invalid tag version response: {}", body).into(), + request_id, + status_code: None, + }) + } + Err(err) => { + let status_code = err.status(); + Err(Error::Http { + source: Box::new(err), + request_id, + status_code, + }) + } + } + } + + /// Resolve a tag to its `(branch, version)` coordinate via the `tags/version` + /// endpoint, since the `/branches/create` contract accepts no `from_tag`. + pub(super) async fn resolve_tag_ref(&self, tag: &str) -> Result<(Option, u64)> { + let request = self + .client + .post(&format!("/v1/table/{}/tags/version/", self.identifier)) + .json(&serde_json::json!({ "tag": tag })); + let (request_id, response) = self.send_unfenced(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let value: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse tag version: {}", e).into(), + request_id: request_id.clone(), + status_code: None, + })?; + let version = value + .get("version") + .and_then(|v| v.as_u64()) + .ok_or_else(|| Error::Http { + source: format!("Invalid tag version response: {}", body).into(), + request_id, + status_code: None, + })?; + let branch = value + .get("branch") + .and_then(|v| v.as_str()) + .map(String::from); + Ok((normalize_branch(branch), version)) + } + + pub(super) async fn describe_with_request( + &self, + request: RequestBuilder, + version: Option, + freshness_request: Option, + ) -> Result { + let mut body = serde_json::json!({ "version": version }); + self.apply_branch_body(&mut body); + let request = request.json(&body); + + let (request_id, response) = if let Some(freshness_request) = freshness_request { + self.send_with_freshness(request, true, freshness_request) + .await? + } else { + self.send_unfenced(request, true).await? + }; + + let response = self.check_table_response(&request_id, response).await?; + + match response.text().await { + Ok(body) => { + let description: TableDescription = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse table description: {}", e).into(), + request_id, + status_code: None, + })?; + if let Some(freshness_request) = freshness_request { + freshness_request.observe_version(&self.freshness, description.version); + } + Ok(description) + } + Err(err) => { + let status_code = err.status(); + Err(Error::Http { + source: Box::new(err), + request_id, + status_code, + }) + } + } + } + + pub(super) async fn send( + &self, + req: RequestBuilder, + with_retry: bool, + ) -> Result<(String, Response)> { + let freshness_request = self.snapshot_freshness_headers(); + self.send_with_freshness(req, with_retry, freshness_request) + .await + } + + pub(super) async fn send_with_freshness( + &self, + req: RequestBuilder, + with_retry: bool, + freshness_request: FreshnessHeaders, + ) -> Result<(String, Response)> { + let req = freshness_request.apply(req); + let res = if with_retry { + self.client.send_with_retry(req, None, true).await? + } else { + self.client.send(req).await? + }; + if res.1.status().is_success() { + freshness_request.observe_headers(&self.freshness, res.1.headers()); + } + Ok(res) + } + + pub(super) async fn send_unfenced( + &self, + req: RequestBuilder, + with_retry: bool, + ) -> Result<(String, Response)> { + if with_retry { + self.client.send_with_retry(req, None, true).await + } else { + self.client.send(req).await + } + } + + pub(in crate::remote) async fn handle_table_not_found( + table_name: &str, + response: reqwest::Response, + request_id: &str, + ) -> Result { + let status = response.status(); + if status == StatusCode::NOT_FOUND { + let body = response.text().await.ok().unwrap_or_default(); + let request_error = Error::Http { + source: body.into(), + request_id: request_id.into(), + status_code: Some(status), + }; + return Err(Error::TableNotFound { + name: table_name.to_string(), + source: Box::new(request_error), + }); + } + Ok(response) + } + + /// Check if a status code should trigger schema cache invalidation + pub(super) fn should_invalidate_cache_for_status(status: StatusCode) -> bool { + // Only invalidate for errors that could be schema-related + // Don't invalidate for auth errors (401, 403) or temporary failures (503, 502) + matches!( + status, + StatusCode::BAD_REQUEST // 400 - could be schema mismatch + | StatusCode::NOT_FOUND // 404 - table might have been recreated + | StatusCode::UNPROCESSABLE_ENTITY // 422 - schema validation error + | StatusCode::INTERNAL_SERVER_ERROR // 500 - could be schema issue on server + ) + } + + pub(super) async fn check_table_response( + &self, + request_id: &str, + response: reqwest::Response, + ) -> Result { + let status = response.status(); + let not_found_result = Self::handle_table_not_found(&self.name, response, request_id).await; + + // Check if we should invalidate cache for 404 errors + if not_found_result.is_err() && Self::should_invalidate_cache_for_status(status) { + self.invalidate_schema_cache(); + } + + let response = not_found_result?; + let result = self.client.check_response(request_id, response).await; + + // Invalidate schema cache on errors that could be schema-related + if result.is_err() && Self::should_invalidate_cache_for_status(status) { + self.invalidate_schema_cache(); + } + + result + } + + pub(super) async fn read_arrow_response( + &self, + request_id: &str, + response: reqwest::Response, + ) -> Result { + let response = self.check_table_response(request_id, response).await?; + + // The header has to be read before the body, which consumes the response. + let content_type = response + .headers() + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let framing = resolve_arrow_ipc_framing(content_type.as_deref(), request_id)?; + + // Buffer the whole body. File framing keeps its footer at the end, so /query + // cannot decode incrementally. Stream framing could, via + // arrow_ipc::reader::StreamDecoder, but fetch_blobs concatenates every batch + // before returning, so no caller would see data sooner. + let body = response.bytes().await.err_to_http(request_id.into())?; + type IpcBatchIterator = + Box> + Send>; + let (schema, batches): (SchemaRef, IpcBatchIterator) = match framing { + ArrowIpcFraming::Stream => { + let reader = StreamReader::try_new(Cursor::new(body), None)?; + (reader.schema(), Box::new(reader)) + } + ArrowIpcFraming::File => { + let reader = FileReader::try_new(Cursor::new(body), None)?; + (reader.schema(), Box::new(reader)) + } + }; + let stream = futures::stream::iter(batches).map_err(DataFusionError::from); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + pub(super) fn apply_query_params( + &self, + body: &mut serde_json::Value, + params: &QueryRequest, + ) -> Result<()> { + params.check_filter()?; + body["prefilter"] = params.prefilter.into(); + // Only forward use_lsm when explicitly set; a server that predates it + // ignores the field and routes as it would by default. + if let Some(use_lsm) = params.use_lsm { + body["use_lsm"] = serde_json::Value::Bool(use_lsm); + } + if let Some(offset) = params.offset { + body["offset"] = serde_json::Value::Number(serde_json::Number::from(offset)); + } + + // Server requires k. + // use isize::MAX as usize to avoid overflow: https://github.com/lancedb/lancedb/issues/2211 + let limit = params.limit.unwrap_or(isize::MAX as usize); + body["k"] = serde_json::Value::Number(serde_json::Number::from(limit)); + + if let Some(filter) = ¶ms.filter { + let filter_sql = match filter { + QueryFilter::Sql(sql) => sql.clone(), + QueryFilter::Datafusion(expr) => expr_to_sql_string(expr)?, + QueryFilter::Substrait(_) => { + return Err(Error::NotSupported { + message: "Substrait filters are not supported for remote queries" + .to_string(), + }); + } + }; + body["filter"] = serde_json::Value::String(filter_sql); + } + + match ¶ms.select { + Select::All => {} + Select::Columns(columns) => { + body["columns"] = serde_json::Value::Array( + columns + .iter() + .map(|s| serde_json::Value::String(s.clone())) + .collect(), + ); + } + Select::Dynamic(pairs) => { + let alias_map = + serde_json::Map::from_iter(pairs.iter().map(|(name, expr)| { + (name.clone(), serde_json::Value::String(expr.clone())) + })); + body["columns"] = alias_map.into(); + } + Select::Expr(pairs) => { + let alias_map: Result> = pairs + .iter() + .map(|(name, expr)| { + expr_to_sql_string(expr) + .map(|sql| (name.clone(), serde_json::Value::String(sql))) + }) + .collect(); + body["columns"] = alias_map?.into(); + } + } + + if params.fast_search { + body["fast_search"] = serde_json::Value::Bool(true); + } + + if params.with_row_id { + body["with_row_id"] = serde_json::Value::Bool(true); + } + + if let Some(full_text_search) = ¶ms.full_text_search { + if full_text_search.wand_factor.is_some() { + return Err(Error::NotSupported { + message: "Wand factor is not yet supported in LanceDB Cloud".into(), + }); + } + + let requires_document_granularity_support = + fts_query_requires_document_granularity_support(&full_text_search.query); + if requires_document_granularity_support + && !self.server_version.support_fts_document_granularity() + { + return Err(Error::NotSupported { + message: + "FTS document granularity requires remote server version 0.6.0 or later" + .into(), + }); + } + + if self.server_version.support_structural_fts() { + body["full_text_query"] = serde_json::json!({ + "query": full_text_search.query.clone(), + }); + } else { + body["full_text_query"] = serde_json::json!({ + "columns": full_text_search.columns().into_iter().collect::>(), + "query": full_text_search.query.query(), + }) + } + } + + if let Some(order_by) = ¶ms.order_by { + body["order_by"] = serde_json::Value::Array( + order_by + .iter() + .map(|o| { + serde_json::json!({ + "column_name": o.column_name, + "ascending": o.ascending, + "nulls_first": o.nulls_first, + }) + }) + .collect(), + ); + } + + Ok(()) + } + + pub(super) fn apply_vector_query_params( + &self, + mut body: serde_json::Value, + query: &VectorQueryRequest, + ) -> Result> { + self.apply_query_params(&mut body, &query.base)?; + + // Apply general parameters, before we dispatch based on number of query vectors. + if let Some(distance_type) = query.distance_type { + body["distance_type"] = serde_json::json!(distance_type); + } + if let Some(approx_mode) = query.approx_mode { + body["approx_mode"] = serde_json::json!(approx_mode); + } + // In 0.23.1 we migrated from `nprobes` to `minimum_nprobes` and `maximum_nprobes`. + // Old client / new server: since minimum_nprobes is missing, fallback to nprobes + // New client / old server: old server will only see nprobes, make sure to set both + // nprobes and minimum_nprobes + // New client / new server: since minimum_nprobes is present, server can ignore nprobes + body["nprobes"] = query.minimum_nprobes.into(); + body["minimum_nprobes"] = query.minimum_nprobes.into(); + if let Some(maximum_nprobes) = query.maximum_nprobes { + body["maximum_nprobes"] = maximum_nprobes.into(); + } else { + body["maximum_nprobes"] = serde_json::Value::Number(Number::from_u128(0).unwrap()) + } + body["lower_bound"] = query.lower_bound.into(); + body["upper_bound"] = query.upper_bound.into(); + body["ef"] = query.ef.into(); + body["refine_factor"] = query.refine_factor.into(); + if let Some(vector_column) = query.column.as_ref() { + body["vector_column"] = serde_json::Value::String(vector_column.clone()); + } + if !query.use_index { + body["bypass_vector_index"] = serde_json::Value::Bool(true); + } + + fn vector_to_json(vector: &arrow_array::ArrayRef) -> Result { + match vector.data_type() { + DataType::Float32 => { + let array = vector + .as_any() + .downcast_ref::() + .unwrap(); + Ok(serde_json::Value::Array( + array + .values() + .iter() + .map(|v| { + serde_json::Number::from_f64(*v as f64) + .map(serde_json::Value::Number) + .ok_or_else(|| Error::InvalidInput { + message: "query vector must contain only finite values" + .into(), + }) + }) + .collect::>>()?, + )) + } + _ => Err(Error::InvalidInput { + message: "VectorQuery vector must be of type Float32".into(), + }), + } + } + + let bodies = match query.query_vector.len() { + 0 => { + // Server takes empty vector, not null or undefined. + body["vector"] = serde_json::Value::Array(Vec::new()); + vec![body] + } + 1 => { + body["vector"] = vector_to_json(&query.query_vector[0])?; + vec![body] + } + _ => { + if self.server_version.support_multivector() { + let vectors = query + .query_vector + .iter() + .map(vector_to_json) + .collect::>>()?; + body["vector"] = serde_json::Value::Array(vectors); + vec![body] + } else { + // Server does not support multiple vectors in a single query. + // We need to send multiple requests. + let mut bodies = Vec::with_capacity(query.query_vector.len()); + for vector in &query.query_vector { + let mut body = body.clone(); + body["vector"] = vector_to_json(vector)?; + bodies.push(body); + } + bodies + } + } + }; + + Ok(bodies) + } + + pub(super) async fn create_multipart_write(&self) -> Result { + let request = self.apply_branch_query(self.client.post(&format!( + "/v1/table/{}/multipart_write/create", + self.identifier + ))); + let (request_id, response) = self.send(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse multipart create response: {}", e).into(), + request_id, + status_code: None, + })?; + parsed["upload_id"] + .as_str() + .map(|s| s.to_string()) + .ok_or_else(|| Error::Http { + source: "Missing upload_id in multipart create response".into(), + request_id: String::new(), + status_code: None, + }) + } + + pub(super) async fn complete_multipart_write(&self, upload_id: &str) -> Result { + let request = self.apply_branch_query( + self.client + .post(&format!( + "/v1/table/{}/multipart_write/complete", + self.identifier + )) + .query(&[("upload_id", upload_id)]), + ); + let (request_id, response) = self.send(request, true).await?; + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse multipart complete response: {}", e).into(), + request_id, + status_code: None, + })?; + let version = parsed["version"].as_u64().ok_or_else(|| Error::Http { + source: "Missing version in multipart complete response".into(), + request_id: String::new(), + status_code: None, + })?; + Ok(AddResult { version }) + } + + pub(super) async fn abort_multipart_write(&self, upload_id: &str) -> Result<()> { + let request = self.apply_branch_query( + self.client + .post(&format!( + "/v1/table/{}/multipart_write/abort", + self.identifier + )) + .query(&[("upload_id", upload_id)]), + ); + let (request_id, response) = self.send(request, true).await?; + self.check_table_response(&request_id, response).await?; + Ok(()) + } + + pub(super) async fn check_mutable(&self) -> Result<()> { + let read_guard = self.version.read().await; + match *read_guard { + None => Ok(()), + Some(version) => Err(Error::NotSupported { + message: format!( + "Cannot mutate table reference fixed at version {}. Call checkout_latest() to get a mutable table reference.", + version + ), + }), + } + } + + pub(super) async fn snapshot_read_state(&self) -> ReadSnapshot { + let version = self.version.read().await; + let (freshness_state, freshness) = + freshness_state_snapshot(&self.freshness, self.client.read_consistency_interval); + ReadSnapshot { + version: *version, + freshness_state, + freshness, + } + } + + /// Snapshot the freshness headers to attach to a single table request. + /// Computed at call time so that retries reuse the same snapshot. + pub(super) fn snapshot_freshness_headers(&self) -> FreshnessHeaders { + freshness_headers_snapshot(&self.freshness, self.client.read_consistency_interval) + } + + pub(super) fn reset_freshness(&self, checkout_baseline: Option, pinned: bool) { + let mut state = self.freshness.lock().unwrap(); + let generation = state.generation.wrapping_add(1); + *state = FreshnessState { + generation, + pinned, + checkout_baseline, + ..FreshnessState::default() + }; + } + + /// Send an LSM operator request with the transport retry layer **off**. + /// + /// Retry policy on these routes belongs to the checkpoint loop, which + /// reads the status and can tell contention from a lost claim. Leaving the + /// transport layer on would re-ask on its own schedule first, and surface + /// an `Error::Retry` whose status the loop would then have to unwrap. + pub(super) async fn send_lsm_route( + &self, + request: RequestBuilder, + ) -> Result<(String, reqwest::Response)> { + let (request_id, response) = self.send(request, false).await?; + let response = self.check_table_response(&request_id, response).await?; + Ok((request_id, response)) + } + + /// Record a version returned by a write so subsequent reads can request at + /// least that version via `x-lancedb-min-version`. A returned `0` from a + /// backward-compatible old server is ignored. + pub(super) fn track_write_version(&self, freshness_request: FreshnessHeaders, version: u64) { + if version == 0 { + return; + } + freshness_request.update_if_current(&self.freshness, |state| { + state.min_version = Some(state.min_version.map_or(version, |v| v.max(version))); + }); + } + + /// Record a committed dataset version observed in a table response so + /// subsequent requests ask for at least this version via + /// `x-lancedb-min-read-version`, + /// giving monotonic reads across load-balanced query nodes. A returned `0` + /// (or absent header from an old server) is ignored. + pub(super) fn track_read_version(&self, version: u64) { + track_read_version(&self.freshness, version); + } + + pub(super) async fn execute_query( + &self, + query: &AnyQuery, + options: &QueryExecutionOptions, + ) -> Result>>> { + let mut request = self + .client + .post(&format!("/v1/table/{}/query/", self.identifier)); + + if let Some(timeout) = options.timeout { + // Also send to server, so it can abort the query if it takes too long. + // (If it doesn't fit into u64, it's not worth sending anyways.) + if let Ok(timeout_ms) = u64::try_from(timeout.as_millis()) { + request = request.header(REQUEST_TIMEOUT_HEADER, timeout_ms); + } + } + + let read_snapshot = self.snapshot_read_state().await; + let query_bodies = self.prepare_query_bodies(query, read_snapshot.version)?; + let requests: Vec = query_bodies + .into_iter() + .map(|body| request.try_clone().unwrap().json(&body)) + .collect(); + + let futures = requests.into_iter().map(|req| async move { + let (request_id, response) = self + .send_with_freshness(req, true, read_snapshot.freshness) + .await?; + self.read_arrow_response(&request_id, response).await + }); + let streams = futures::future::try_join_all(futures); + + if let Some(timeout) = options.timeout { + let timeout_future = tokio::time::sleep(timeout); + tokio::pin!(timeout_future); + tokio::pin!(streams); + tokio::select! { + _ = &mut timeout_future => { + Err(Error::Other { + message: format!("Query timeout after {} ms", timeout.as_millis()), + source: None, + }) + } + result = &mut streams => { + Ok(result?) + } + } + } else { + Ok(streams.await?) + } + } + + pub(super) fn prepare_query_bodies( + &self, + query: &AnyQuery, + version: Option, + ) -> Result> { + let query = query.canonicalized()?; + let mut base_body = serde_json::json!({ "version": version }); + self.apply_branch_body(&mut base_body); + + match &query { + AnyQuery::Query(query) => { + let mut body = base_body.clone(); + self.apply_query_params(&mut body, query)?; + // Empty vector can be passed if no vector search is performed. + body["vector"] = serde_json::Value::Array(Vec::new()); + Ok(vec![body]) + } + AnyQuery::VectorQuery(query) => self.apply_vector_query_params(base_body, query), + } + } + + pub(super) fn invalidate_schema_cache(&self) { + self.schema_cache.invalidate(); + } + + pub(super) fn handle_error_invalidation(&self, error: &Error) { + let status_code = match error { + Error::Http { status_code, .. } => *status_code, + Error::Retry { status_code, .. } => *status_code, + _ => None, + }; + if let Some(status_code) = status_code + && Self::should_invalidate_cache_for_status(status_code) + { + self.invalidate_schema_cache(); + } + } +} diff --git a/rust/lancedb/src/remote/table/tags.rs b/rust/lancedb/src/remote/table/tags.rs new file mode 100644 index 000000000..002ec352e --- /dev/null +++ b/rust/lancedb/src/remote/table/tags.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +pub struct RemoteTags<'a, S: HttpSend = Sender> { + pub(super) inner: &'a RemoteTable, +} + +#[async_trait] +impl Tags for RemoteTags<'_, S> { + async fn list(&self) -> Result> { + let request = self + .inner + .client + .post(&format!("/v1/table/{}/tags/list/", self.inner.identifier)); + let (request_id, response) = self.inner.send_unfenced(request, true).await?; + let response = self + .inner + .check_table_response(&request_id, response) + .await?; + + match response.text().await { + Ok(body) => { + // Explicitly tell serde_json what type we want to deserialize into + let tags_map: HashMap = + serde_json::from_str(&body).map_err(|e| Error::Http { + source: format!("Failed to parse tags list: {}", e).into(), + request_id, + status_code: None, + })?; + + Ok(tags_map) + } + Err(err) => { + let status_code = err.status(); + Err(Error::Http { + source: Box::new(err), + request_id, + status_code, + }) + } + } + } + + async fn get_version(&self, tag: &str) -> Result { + let request = self.inner.client.post(&format!( + "/v1/table/{}/tags/version/", + self.inner.identifier + )); + self.inner + .resolve_tag_version_with_request(tag, request, false) + .await + } + + async fn create(&mut self, tag: &str, version: u64) -> Result<()> { + let mut body = serde_json::json!({ + "tag": tag, + "version": version + }); + self.inner.apply_branch_body(&mut body); + let request = self + .inner + .client + .post(&format!("/v1/table/{}/tags/create/", self.inner.identifier)) + .json(&body); + + let (request_id, response) = self.inner.send(request, true).await?; + self.inner + .check_table_response(&request_id, response) + .await?; + Ok(()) + } + + async fn delete(&mut self, tag: &str) -> Result<()> { + let request = self + .inner + .client + .post(&format!("/v1/table/{}/tags/delete/", self.inner.identifier)) + .json(&serde_json::json!({ "tag": tag })); + + let (request_id, response) = self.inner.send_unfenced(request, true).await?; + self.inner + .check_table_response(&request_id, response) + .await?; + Ok(()) + } + + async fn update(&mut self, tag: &str, version: u64) -> Result<()> { + let mut body = serde_json::json!({ + "tag": tag, + "version": version + }); + self.inner.apply_branch_body(&mut body); + let request = self + .inner + .client + .post(&format!("/v1/table/{}/tags/update/", self.inner.identifier)) + .json(&body); + + let (request_id, response) = self.inner.send(request, true).await?; + self.inner + .check_table_response(&request_id, response) + .await?; + Ok(()) + } +} diff --git a/rust/lancedb/src/remote/table/tests/index.rs b/rust/lancedb/src/remote/table/tests/index.rs index 3b40ab454..27f7cbc7b 100644 --- a/rust/lancedb/src/remote/table/tests/index.rs +++ b/rust/lancedb/src/remote/table/tests/index.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +use super::super::indices::deserialize_created_at; use super::*; #[tokio::test] diff --git a/rust/lancedb/src/remote/table/write.rs b/rust/lancedb/src/remote/table/write.rs new file mode 100644 index 000000000..1106a741d --- /dev/null +++ b/rust/lancedb/src/remote/table/write.rs @@ -0,0 +1,219 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::*; + +impl RemoteTable { + pub(super) fn is_retryable_write_error(&self, err: &Error) -> bool { + match err { + Error::Http { + source, + status_code, + .. + } => { + // Don't retry read errors (is_body/is_decode): the + // server may have committed the write already, and + // without an idempotency key we'd duplicate data. + source + .downcast_ref::() + .is_some_and(|e| e.is_connect()) + || status_code.is_some_and(|s| self.client.retry_config.statuses.contains(&s)) + } + // send_with_retry exhausted its internal retries on a retryable + // status. The outer loop can still retry the whole operation with + // a fresh session. + Error::Retry { status_code, .. } => { + status_code.is_some_and(|s| self.client.retry_config.statuses.contains(&s)) + } + _ => false, + } + } + + pub(super) async fn add_single_partition( + &self, + output: PreprocessingOutput, + ) -> Result { + use crate::remote::retry::RetryCounter; + + let _guard = output.tracker.as_ref().map(|t| t.track_task()); + let freshness_request = self.snapshot_freshness_headers(); + + let mut insert: Arc = Arc::new( + RemoteWriteExec::new( + self.name.clone(), + self.identifier.clone(), + self.client.clone(), + output.plan, + WriteOp::Insert { + overwrite: output.overwrite, + }, + output.tracker.clone(), + self.branch.clone(), + ) + .with_freshness( + self.freshness.clone(), + self.client.read_consistency_interval, + ), + ); + + let mut retry_counter = + RetryCounter::new(&self.client.retry_config, uuid::Uuid::new_v4().to_string()); + + loop { + let stream = execute_plan(insert.clone(), Default::default())?; + let result: Result> = stream.try_collect().await.map_err(Error::from); + + match result { + Ok(_) => { + let add_result = (insert.as_ref() as &dyn std::any::Any) + .downcast_ref::>() + .and_then(|i| i.add_result()) + .unwrap_or(AddResult { version: 0 }); + + if output.overwrite { + self.invalidate_schema_cache(); + } + self.track_write_version(freshness_request, add_result.version); + + return Ok(add_result); + } + Err(err) if output.rescannable && self.is_retryable_write_error(&err) => { + retry_counter.increment_from_error(err)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + insert = insert.reset_state()?; + continue; + } + Err(err) => return Err(err), + } + } + } + + pub(super) async fn add_multipart( + &self, + output: PreprocessingOutput, + num_partitions: usize, + ) -> Result { + use crate::remote::retry::RetryCounter; + + let mut retry_counter = + RetryCounter::new(&self.client.retry_config, uuid::Uuid::new_v4().to_string()); + + loop { + let freshness_request = self.snapshot_freshness_headers(); + let upload_id = self.create_multipart_write().await?; + + let result = self + .execute_multipart_inserts(&upload_id, &output, num_partitions) + .await; + + match result { + Ok(()) => match self.complete_multipart_write(&upload_id).await { + Ok(result) => { + if output.overwrite { + self.invalidate_schema_cache(); + } + self.track_write_version(freshness_request, result.version); + return Ok(result); + } + Err(e) => { + if let Err(abort_err) = self.abort_multipart_write(&upload_id).await { + log::warn!( + "Failed to abort multipart write {}: {}", + upload_id, + abort_err + ); + } + if output.rescannable && self.is_retryable_write_error(&e) { + retry_counter.increment_from_error(e)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + continue; + } + return Err(e); + } + }, + Err(e) => { + if let Err(abort_err) = self.abort_multipart_write(&upload_id).await { + log::warn!( + "Failed to abort multipart write {}: {}", + upload_id, + abort_err + ); + } + if output.rescannable && self.is_retryable_write_error(&e) { + retry_counter.increment_from_error(e)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + continue; + } + return Err(e); + } + } + } + } + + pub(super) async fn execute_multipart_inserts( + &self, + upload_id: &str, + output: &PreprocessingOutput, + num_partitions: usize, + ) -> Result<()> { + debug_assert!( + output.rescannable, + "multipart inserts require rescannable input for retry support" + ); + + let plan = Arc::new( + datafusion_physical_plan::repartition::RepartitionExec::try_new( + output.plan.clone(), + datafusion_physical_plan::Partitioning::RoundRobinBatch(num_partitions), + )?, + ) as Arc; + + let insert = Arc::new( + RemoteWriteExec::new_multipart( + self.name.clone(), + self.identifier.clone(), + self.client.clone(), + plan, + output.overwrite, + upload_id.to_string(), + output.tracker.clone(), + self.branch.clone(), + self.client.max_bytes_per_request(), + self.client.max_request_duration(), + ) + .with_freshness( + self.freshness.clone(), + self.client.read_consistency_interval, + ), + ); + + let task_ctx = Arc::new(datafusion_execution::TaskContext::default()); + let tracker = output.tracker.clone(); + let mut join_set = tokio::task::JoinSet::new(); + for partition in 0..num_partitions { + let exec = insert.clone(); + let ctx = task_ctx.clone(); + let tracker = tracker.clone(); + join_set.spawn(async move { + let _guard = tracker.as_ref().map(|t| t.track_task()); + let mut stream = exec + .execute(partition, ctx) + .map_err(|e| -> Error { e.into() })?; + while let Some(batch) = stream.next().await { + batch.map_err(|e| -> Error { e.into() })?; + } + Ok::<_, Error>(()) + }); + } + + // JoinSet aborts all remaining tasks when dropped, so if we return + // early on error the orphaned tasks are automatically cancelled. + while let Some(result) = join_set.join_next().await { + result.map_err(|e| Error::Runtime { + message: format!("Insert task panicked: {}", e), + })??; + } + + Ok(()) + } +}