diff --git a/nodejs/src/remote.rs b/nodejs/src/remote.rs index 8afbfd925..4bdb5685e 100644 --- a/nodejs/src/remote.rs +++ b/nodejs/src/remote.rs @@ -232,6 +232,10 @@ impl From for lancedb::remote::ClientConfig { tls_config: config.tls_config.map(Into::into), header_provider: None, // the header provider is set separately later user_id: config.user_id, + // Resolved from LANCE_CLIENT_MAX_BYTES_PER_REQUEST or the default. + max_bytes_per_request: None, + // Resolved from LANCE_CLIENT_MAX_REQUEST_DURATION or the read timeout. + max_request_duration: None, } } } diff --git a/python/src/connection.rs b/python/src/connection.rs index 16bb71f80..e22813640 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -800,6 +800,10 @@ impl From for lancedb::remote::ClientConfig { tls_config: value.tls_config.map(Into::into), header_provider, user_id: value.user_id, + // Resolved from LANCE_CLIENT_MAX_BYTES_PER_REQUEST or the default. + max_bytes_per_request: None, + // Resolved from LANCE_CLIENT_MAX_REQUEST_DURATION or the read timeout. + max_request_duration: None, } } } diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index d4a65fe15..2f4ac03d3 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -47,6 +47,20 @@ pub trait HeaderProvider: Send + Sync + std::fmt::Debug { async fn get_headers(&self) -> Result>; } +/// Default maximum bytes per insert request (8 GiB). +/// +/// Sized so a multipart part can hold at least one full Lance data file (the +/// default is 1M rows / 90 GB per file), which keeps fragments from being split +/// into undersized files across parts. The time-based cut +/// ([`DEFAULT_MAX_REQUEST_DURATION_DIVISOR`]) bounds request duration on slow +/// uploads, so a large byte budget does not risk the read timeout. +const DEFAULT_MAX_BYTES_PER_REQUEST: u64 = 8 * 1024 * 1024 * 1024; + +/// The default max request duration is the read timeout divided by this, leaving +/// headroom for the server to finalize and acknowledge a part before the read +/// timeout (which also covers the request-body upload) fires. +const DEFAULT_MAX_REQUEST_DURATION_DIVISOR: u32 = 2; + /// Configuration for the LanceDB Cloud HTTP client. #[derive(Clone)] pub struct ClientConfig { @@ -71,6 +85,33 @@ pub struct ClientConfig { /// Alternatively, set `LANCEDB_USER_ID_ENV_KEY` to specify another environment /// variable that contains the user ID value. pub user_id: Option, + /// Maximum number of bytes to send in a single insert HTTP request. + /// + /// During a multipart write, each partition's data is split into one or more + /// parts of at most this many (Arrow IPC, compressed) bytes, each uploaded as + /// a separate request under the shared upload id. This bounds how long any + /// one request stays open, so large bulk ingests do not exceed the client + /// read timeout while the server streams the part to object storage. + /// + /// The request body is still streamed (not buffered), so this does not + /// increase peak memory. Set to `Some(0)` to disable splitting (one request + /// per partition). You can also set the `LANCE_CLIENT_MAX_BYTES_PER_REQUEST` + /// environment variable. Defaults to 8 GiB. + pub max_bytes_per_request: Option, + /// Maximum wall-clock time to spend uploading a single insert HTTP request. + /// + /// Complements [`Self::max_bytes_per_request`]: during a multipart write a + /// part is cut when it reaches either the byte budget or this duration, + /// whichever comes first. The client read timeout also covers the + /// request-body upload, so a slow or throttled upload of a large part can + /// hit that timeout before the byte budget is reached; cutting by time keeps + /// each request short enough that it completes (and the server acknowledges + /// the part) within the read timeout. + /// + /// Set to `Some(Duration::ZERO)` to disable the time-based cut. You can also + /// set the `LANCE_CLIENT_MAX_REQUEST_DURATION` environment variable (integer + /// seconds). Defaults to half the resolved read timeout. + pub max_request_duration: Option, } impl std::fmt::Debug for ClientConfig { @@ -87,6 +128,8 @@ impl std::fmt::Debug for ClientConfig { &self.header_provider.as_ref().map(|_| "Some(...)"), ) .field("user_id", &self.user_id) + .field("max_bytes_per_request", &self.max_bytes_per_request) + .field("max_request_duration", &self.max_request_duration) .finish() } } @@ -102,6 +145,8 @@ impl Default for ClientConfig { tls_config: None, header_provider: None, user_id: None, + max_bytes_per_request: None, + max_request_duration: None, } } } @@ -248,6 +293,16 @@ pub struct RestfulLanceDbClient { /// Connection-level read consistency interval. Drives the /// `x-lancedb-min-timestamp` freshness header sent on read requests. pub(crate) read_consistency_interval: Option, + // Note the `Option` here means the opposite of the same-named + // `ClientConfig` fields: those are pre-resolution, where `None` means "fall + // back to env var / default". These are post-resolution (see + // `resolve_max_bytes_per_request` / `resolve_max_request_duration`), where a + // default has already been applied and `None` means the feature is disabled. + /// Maximum bytes per insert request. `None` disables request splitting. + pub(crate) max_bytes_per_request: Option, + /// Maximum wall-clock time per insert request. `None` disables the + /// time-based part cut. + pub(crate) max_request_duration: Option, } impl std::fmt::Debug for RestfulLanceDbClient { @@ -429,6 +484,10 @@ impl RestfulLanceDbClient { }; debug!("Created client for host: {}", host); let retry_config = client_config.retry_config.clone().try_into()?; + let max_bytes_per_request = + Self::resolve_max_bytes_per_request(client_config.max_bytes_per_request)?; + let max_request_duration = + Self::resolve_max_request_duration(client_config.max_request_duration, read_timeout)?; Ok(Self { client, host, @@ -440,8 +499,52 @@ impl RestfulLanceDbClient { .unwrap_or("$".to_string()), header_provider: client_config.header_provider, read_consistency_interval, + max_bytes_per_request, + max_request_duration, }) } + + /// Resolve the max bytes per insert request from config, environment, or the + /// default. A value of `0` (from either source) disables request splitting. + fn resolve_max_bytes_per_request(passed: Option) -> Result> { + let value = if let Some(value) = passed { + value + } else if let Ok(env) = std::env::var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST") { + env.parse::().map_err(|_| Error::InvalidInput { + message: format!( + "LANCE_CLIENT_MAX_BYTES_PER_REQUEST must be a non-negative integer, got '{}'", + env + ), + })? + } else { + DEFAULT_MAX_BYTES_PER_REQUEST + }; + Ok((value > 0).then_some(value)) + } + + /// Resolve the max request duration from config, environment, or a default + /// derived from the read timeout. A zero duration (from either source) + /// disables the time-based cut. + fn resolve_max_request_duration( + passed: Option, + read_timeout: Duration, + ) -> Result> { + let value = if let Some(value) = passed { + value + } else if let Ok(env) = std::env::var("LANCE_CLIENT_MAX_REQUEST_DURATION") { + let secs = env.parse::().map_err(|_| Error::InvalidInput { + message: format!( + "LANCE_CLIENT_MAX_REQUEST_DURATION must be a non-negative integer \ + number of seconds, got '{}'", + env + ), + })?; + Duration::from_secs(secs) + } else { + read_timeout / DEFAULT_MAX_REQUEST_DURATION_DIVISOR + }; + Ok((!value.is_zero()).then_some(value)) + } } impl RestfulLanceDbClient { @@ -449,6 +552,18 @@ impl RestfulLanceDbClient { &self.host } + /// Maximum bytes per insert request, or `None` if request splitting is + /// disabled. + pub(crate) fn max_bytes_per_request(&self) -> Option { + self.max_bytes_per_request + } + + /// Maximum wall-clock time per insert request, or `None` if the time-based + /// cut is disabled. + pub(crate) fn max_request_duration(&self) -> Option { + self.max_request_duration + } + pub fn default_headers( api_key: &str, region: &str, @@ -875,6 +990,8 @@ pub mod test_utils { id_delimiter: "$".to_string(), header_provider: None, read_consistency_interval, + max_bytes_per_request: None, + max_request_duration: None, } } @@ -900,6 +1017,12 @@ pub mod test_utils { id_delimiter: config.id_delimiter.unwrap_or_else(|| "$".to_string()), header_provider: config.header_provider, read_consistency_interval: None, + max_bytes_per_request: config + .max_bytes_per_request + .and_then(|v| (v > 0).then_some(v)), + max_request_duration: config + .max_request_duration + .and_then(|v| (!v.is_zero()).then_some(v)), } } } @@ -1103,6 +1226,8 @@ mod tests { id_delimiter: "+".to_string(), header_provider: Some(Arc::new(provider) as Arc), read_consistency_interval: None, + max_bytes_per_request: None, + max_request_duration: None, }; // Apply dynamic headers @@ -1139,6 +1264,8 @@ mod tests { id_delimiter: "+".to_string(), header_provider: Some(Arc::new(provider) as Arc), read_consistency_interval: None, + max_bytes_per_request: None, + max_request_duration: None, }; // Apply dynamic headers @@ -1177,6 +1304,8 @@ mod tests { id_delimiter: "+".to_string(), header_provider: Some(Arc::new(provider) as Arc), read_consistency_interval: None, + max_bytes_per_request: None, + max_request_duration: None, }; // Header provider errors should fail the request @@ -1288,4 +1417,193 @@ mod tests { std::env::remove_var("LANCEDB_USER_ID"); } } + + #[test] + fn test_resolve_max_bytes_passed_value_wins() { + // An explicit config value is used verbatim; env/default are not consulted. + let resolved = + RestfulLanceDbClient::::resolve_max_bytes_per_request(Some(1234)).unwrap(); + assert_eq!(resolved, Some(1234)); + } + + #[test] + fn test_resolve_max_bytes_zero_disables() { + let resolved = + RestfulLanceDbClient::::resolve_max_bytes_per_request(Some(0)).unwrap(); + assert_eq!(resolved, None); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_bytes_default_when_unset() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST"); + } + let resolved = RestfulLanceDbClient::::resolve_max_bytes_per_request(None).unwrap(); + assert_eq!(resolved, Some(DEFAULT_MAX_BYTES_PER_REQUEST)); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_bytes_from_env() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096"); + } + let resolved = RestfulLanceDbClient::::resolve_max_bytes_per_request(None).unwrap(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST"); + } + assert_eq!(resolved, Some(4096)); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_bytes_env_zero_disables() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "0"); + } + let resolved = RestfulLanceDbClient::::resolve_max_bytes_per_request(None).unwrap(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST"); + } + assert_eq!(resolved, None); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_bytes_config_overrides_env() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "4096"); + } + // A config value takes precedence over the environment variable. + let resolved = + RestfulLanceDbClient::::resolve_max_bytes_per_request(Some(1234)).unwrap(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST"); + } + assert_eq!(resolved, Some(1234)); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_bytes_invalid_env_errors() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST", "not-a-number"); + } + let err = RestfulLanceDbClient::::resolve_max_bytes_per_request(None).unwrap_err(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_BYTES_PER_REQUEST"); + } + assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}"); + } + + #[test] + fn test_resolve_max_request_duration_passed_value_wins() { + let resolved = RestfulLanceDbClient::::resolve_max_request_duration( + Some(Duration::from_secs(42)), + Duration::from_secs(300), + ) + .unwrap(); + assert_eq!(resolved, Some(Duration::from_secs(42))); + } + + #[test] + fn test_resolve_max_request_duration_zero_disables() { + let resolved = RestfulLanceDbClient::::resolve_max_request_duration( + Some(Duration::ZERO), + Duration::from_secs(300), + ) + .unwrap(); + assert_eq!(resolved, None); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_request_duration_default_is_half_read_timeout() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION"); + } + let resolved = RestfulLanceDbClient::::resolve_max_request_duration( + None, + Duration::from_secs(300), + ) + .unwrap(); + assert_eq!(resolved, Some(Duration::from_secs(150))); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_request_duration_from_env_seconds() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "30"); + } + let resolved = RestfulLanceDbClient::::resolve_max_request_duration( + None, + Duration::from_secs(300), + ) + .unwrap(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION"); + } + assert_eq!(resolved, Some(Duration::from_secs(30))); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_request_duration_env_zero_disables() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "0"); + } + let resolved = RestfulLanceDbClient::::resolve_max_request_duration( + None, + Duration::from_secs(300), + ) + .unwrap(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION"); + } + assert_eq!(resolved, None); + } + + #[test] + #[serial(request_limits_env)] + fn test_resolve_max_request_duration_invalid_env_errors() { + let _guard = lock_env(); + // SAFETY: This is only called in tests + unsafe { + std::env::set_var("LANCE_CLIENT_MAX_REQUEST_DURATION", "12.5"); + } + let err = RestfulLanceDbClient::::resolve_max_request_duration( + None, + Duration::from_secs(300), + ) + .unwrap_err(); + // SAFETY: This is only called in tests + unsafe { + std::env::remove_var("LANCE_CLIENT_MAX_REQUEST_DURATION"); + } + assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}"); + } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 9f28f1657..90e0d3dd6 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1361,6 +1361,8 @@ impl RemoteTable { upload_id.to_string(), output.tracker.clone(), self.branch.clone(), + self.client.max_bytes_per_request(), + self.client.max_request_duration(), )); let task_ctx = Arc::new(datafusion_execution::TaskContext::default()); @@ -1863,25 +1865,33 @@ impl BaseTable for RemoteTable { let table_schema = self.schema().await?; let table_def = TableDefinition::try_from_rich_schema(table_schema.clone())?; - let num_partitions = if let Some(parallelism) = add.write_parallelism { - if parallelism > 1 && self.server_version.support_multipart_write() { - parallelism - } else { - 1 - } - } else if self.server_version.support_multipart_write() { - // Peek at the first batch to estimate write partitions, same as NativeTable. + let num_partitions = if self.server_version.support_multipart_write() { + // Peek at the first batch to estimate write partitions (same as + // NativeTable) and, regardless of `write_parallelism`, to detect a + // fully empty input. A multipart write creates its upload session + // before any partition executes; if the input turns out to have no + // batches at all, no partition ever stages a part (see + // `send_multipart_chunked`), so completing the write has nothing to + // commit and e.g. `mode=overwrite` would be silently dropped. Route + // empty input through the single-request path instead, which always + // sends one schema-only request. let mut peeked = PeekedScannable::new(add.data); - let n = if let Some(first_batch) = peeked.peek().await { - let max_partitions = lance_core::utils::tokio::get_num_compute_intensive_cpus(); - estimate_write_partitions( - first_batch.get_array_memory_size(), - first_batch.num_rows(), - peeked.num_rows(), - max_partitions, - ) - } else { - 1 + let n = match peeked.peek().await { + Some(first_batch) => match add.write_parallelism { + Some(parallelism) if parallelism > 1 => parallelism, + Some(_) => 1, + None => { + let max_partitions = + lance_core::utils::tokio::get_num_compute_intensive_cpus(); + estimate_write_partitions( + first_batch.get_array_memory_size(), + first_batch.num_rows(), + peeked.num_rows(), + max_partitions, + ) + } + }, + None => 1, }; add.data = Box::new(peeked); n @@ -7166,6 +7176,76 @@ mod tests { assert_eq!(insert_count.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn test_multipart_write_empty_overwrite_uses_single_partition() { + // A multipart write creates its upload session before any partition + // executes. If the input has no batches at all, every partition would + // stage nothing (see `send_multipart_chunked`), so completing the write + // would have nothing to commit and `mode=overwrite` would be silently + // dropped. An explicit `write_parallelism` must not force the multipart + // path for empty input; it should fall back to the single-request path, + // which always sends one schema-only request and carries `mode=overwrite`. + let insert_count = Arc::new(AtomicUsize::new(0)); + let multipart_count = Arc::new(AtomicUsize::new(0)); + + let insert_count_c = insert_count.clone(); + let multipart_count_c = multipart_count.clone(); + + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 4, 0), + move |request| { + let path = request.url().path(); + + if path == "/v1/table/my_table/describe/" { + return simple_describe_response(); + } + + if path.contains("multipart_write") { + multipart_count_c.fetch_add(1, Ordering::SeqCst); + panic!("Should not use multipart write endpoints for empty input"); + } + + if path == "/v1/table/my_table/insert/" { + let query = request.url().query().unwrap_or(""); + assert!( + !query.contains("upload_id"), + "Should not have upload_id for empty input" + ); + assert!( + query.contains("mode=overwrite"), + "Should carry mode=overwrite, got query: {}", + query + ); + insert_count_c.fetch_add(1, Ordering::SeqCst); + return http::Response::builder() + .status(200) + .body(r#"{"version": 2}"#.to_string()) + .unwrap(); + } + + panic!("Unexpected request path: {}", path); + }, + ); + + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); + let empty_batches: Vec> = + Vec::new(); + let data: Box = + Box::new(RecordBatchIterator::new(empty_batches, schema)); + let result = table + .add(data) + .mode(AddDataMode::Overwrite) + .write_parallelism(4) + .execute() + .await + .unwrap(); + + assert_eq!(result.version, 2); + assert_eq!(multipart_count.load(Ordering::SeqCst), 0); + assert_eq!(insert_count.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn test_multipart_write_abort_on_insert_failure() { let create_count = Arc::new(AtomicUsize::new(0)); diff --git a/rust/lancedb/src/remote/table/insert.rs b/rust/lancedb/src/remote/table/insert.rs index 649ebc225..4b0aa5480 100644 --- a/rust/lancedb/src/remote/table/insert.rs +++ b/rust/lancedb/src/remote/table/insert.rs @@ -4,6 +4,7 @@ //! DataFusion ExecutionPlan for inserting data into remote LanceDB tables. use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use arrow_array::{ArrayRef, RecordBatch, UInt64Array}; use arrow_ipc::CompressionType; @@ -15,7 +16,7 @@ use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, }; -use futures::StreamExt; +use futures::{SinkExt, StreamExt}; use http::header::CONTENT_TYPE; use lance::io::exec::utils::InstrumentedRecordBatchStreamAdapter; @@ -49,6 +50,15 @@ pub struct RemoteInsertExec { tracker: Option>, /// Branch to write to via `?branch=`. `None` targets the main branch. branch: Option, + /// For multipart writes, split each partition into parts of at most this + /// many bytes, each uploaded as a separate request. `None` sends the whole + /// partition as a single request. + max_bytes_per_request: Option, + /// For multipart writes, also cut a part once it has been uploading for this + /// long, even if it has not reached `max_bytes_per_request`. Bounds request + /// duration on slow/throttled uploads so no request exceeds the read + /// timeout. `None` disables the time-based cut. + max_request_duration: Option, } impl RemoteInsertExec { @@ -63,7 +73,7 @@ impl RemoteInsertExec { branch: Option, ) -> Self { Self::new_inner( - table_name, identifier, client, input, overwrite, None, tracker, branch, + table_name, identifier, client, input, overwrite, None, tracker, branch, None, None, ) } @@ -82,6 +92,8 @@ impl RemoteInsertExec { upload_id: String, tracker: Option>, branch: Option, + max_bytes_per_request: Option, + max_request_duration: Option, ) -> Self { Self::new_inner( table_name, @@ -92,6 +104,8 @@ impl RemoteInsertExec { Some(upload_id), tracker, branch, + max_bytes_per_request, + max_request_duration, ) } @@ -105,6 +119,8 @@ impl RemoteInsertExec { upload_id: Option, tracker: Option>, branch: Option, + max_bytes_per_request: Option, + max_request_duration: Option, ) -> Self { let num_partitions = if upload_id.is_some() { input.output_partitioning().partition_count() @@ -131,6 +147,8 @@ impl RemoteInsertExec { upload_id, tracker, branch, + max_bytes_per_request, + max_request_duration, } } @@ -214,6 +232,238 @@ impl RemoteInsertExec { } } +/// Shared context for the requests of a single partition's multipart upload. +/// These values are identical for every part; only the part id and streamed +/// body differ between requests. Bundling them keeps the per-part helpers from +/// each threading the same handful of arguments. +struct PartRequestCtx<'a, S: HttpSend> { + client: &'a RestfulLanceDbClient, + identifier: &'a str, + table_name: &'a str, + upload_id: &'a str, + branch: Option<&'a str>, + overwrite: bool, +} + +impl PartRequestCtx<'_, S> { + /// Upload a partition as one or more multipart parts, cutting a new part + /// whenever the current one reaches `max_bytes` (Arrow IPC, compressed) or + /// has been uploading for `max_duration`, whichever comes first. + /// + /// Each part is a separate `/insert?upload_id=...&upload_part_id=...` request + /// whose body is still streamed through a bounded channel, so peak memory + /// stays at a couple of batches regardless of `max_bytes`. The server stages + /// every part under the shared `upload_id` and merges them atomically when + /// the caller completes the multipart write. An empty partition stages + /// nothing: the multipart write always has at least one non-empty partition + /// to commit. + /// + /// The byte budget targets a good on-disk fragment size; the duration budget + /// bounds request time so a slow or throttled upload does not keep a request + /// open past the client read timeout (which also covers the request body). + async fn send_multipart_chunked( + &self, + max_bytes: u64, + max_duration: Option, + mut input: SendableRecordBatchStream, + tracker: Option>, + ) -> DataFusionResult<()> { + let schema = input.schema(); + + // A part always starts from a batch we already hold: the first batch of + // the partition, or the look-ahead batch from the previous part. This + // keeps empty partitions from staging a part and stops a size cut that + // lands exactly on the end of input from emitting a trailing empty part. + let mut first = match input.next().await { + Some(batch) => batch?, + None => return Ok(()), + }; + + loop { + let input_ended = self + .send_one_part( + &schema, + max_bytes, + max_duration, + first, + &mut input, + &tracker, + ) + .await?; + + if input_ended { + break; + } + + first = match input.next().await { + Some(batch) => batch?, + None => break, + }; + } + + Ok(()) + } + + /// Build the `/insert` request for a single multipart part. + fn build_part_request(&self, part_id: &str, body: reqwest::Body) -> reqwest::RequestBuilder { + let mut request = self + .client + .post(&format!("/v1/table/{}/insert/", self.identifier)) + .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE) + .query(&[("upload_id", self.upload_id)]) + .query(&[("upload_part_id", part_id)]); + // Every part of an overwrite carries `mode=overwrite`. The server records + // it against the shared `upload_id` and applies the overwrite once, when + // the multipart write is completed, rather than per part. + if self.overwrite { + request = request.query(&[("mode", "overwrite")]); + } + if let Some(b) = self.branch { + request = request.query(&[("branch", b)]); + } + request.body(body) + } + + /// Send a single part's request and drain the response, mapping HTTP and + /// table-not-found errors into `DataFusionError`. + async fn send_part_request(&self, request: reqwest::RequestBuilder) -> DataFusionResult<()> { + let (request_id, response) = self + .client + .send(request) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let response = + RemoteTable::::handle_table_not_found(self.table_name, response, &request_id) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let response = self + .client + .check_response(&request_id, response) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + response.bytes().await.map_err(|e| { + DataFusionError::External(Box::new(Error::Http { + source: Box::new(e), + request_id: request_id.clone(), + status_code: None, + })) + })?; + Ok(()) + } + + /// Stream one part, starting from `first` and pulling from `input` until the + /// part reaches `max_bytes`, has been uploading for `max_duration`, or the + /// input ends. The body is streamed through a bounded channel concurrently + /// with the request, so peak memory stays at a couple of batches. Wire bytes + /// are recorded on `tracker` as each chunk is produced, so progress advances + /// smoothly rather than jumping once per completed part. Returns whether the + /// input was exhausted while filling this part. + async fn send_one_part( + &self, + schema: &arrow_schema::SchemaRef, + max_bytes: u64, + max_duration: Option, + first: RecordBatch, + input: &mut SendableRecordBatchStream, + tracker: &Option>, + ) -> DataFusionResult { + let (mut chunk_tx, chunk_rx) = + futures::channel::mpsc::channel::, std::io::Error>>(2); + let body = reqwest::Body::wrap_stream(chunk_rx); + + let part_id = uuid::Uuid::new_v4().to_string(); + let request = self.build_part_request(&part_id, body); + + // Measured from just before the request is sent, matching the window the + // client read timeout applies to the upload. + let started = Instant::now(); + let tracker = tracker.clone(); + // Unlike `stream_as_http_body`, this producer also cuts the part at the + // byte/time budget and reports back whether the input ended, so it drives + // its own bounded mpsc channel joined with the request instead of reusing + // that helper. + let producer = async move { + let options = arrow_ipc::writer::IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::LZ4_FRAME)) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let mut writer = + arrow_ipc::writer::StreamWriter::try_new_with_options(Vec::new(), schema, options) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + let mut part_bytes: u64 = 0; + let mut input_ended = false; + let mut pending = Some(first); + loop { + let batch = match pending.take() { + Some(batch) => batch, + None => match input.next().await { + Some(Ok(batch)) => batch, + Some(Err(e)) => { + // Abort the body so the server does not treat the + // truncated stream as a successful write; the + // original error is surfaced to the caller. + let _ = chunk_tx + .send(Err(std::io::Error::other("input stream error"))) + .await; + return Err(e); + } + None => { + input_ended = true; + break; + } + }, + }; + writer + .write(&batch) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let chunk = std::mem::take(writer.get_mut()); + let chunk_len = chunk.len(); + part_bytes += chunk_len as u64; + if chunk_tx.send(Ok(chunk)).await.is_err() { + // The request finished or failed; stop producing. + break; + } + if let Some(ref t) = tracker { + t.record_bytes(chunk_len); + } + if part_bytes >= max_bytes + || max_duration.is_some_and(|limit| started.elapsed() >= limit) + { + break; + } + } + + writer + .finish() + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let tail = std::mem::take(writer.get_mut()); + if !tail.is_empty() { + let tail_len = tail.len(); + if chunk_tx.send(Ok(tail)).await.is_ok() + && let Some(ref t) = tracker + { + t.record_bytes(tail_len); + } + } + Ok::(input_ended) + }; + + let send = self.send_part_request(request); + + // `join!` rather than `tokio::spawn`: the producer borrows `input` (and + // `schema`), so it cannot satisfy the `'static` bound a spawned task + // needs. Running both futures on this task lets them make progress + // concurrently without that constraint. + let (producer_result, send_result) = futures::join!(producer, send); + // Prefer the producer error (e.g. NaN rejection) over any HTTP error it + // induced. + let input_ended = producer_result?; + send_result?; + + Ok(input_ended) + } +} + impl DisplayAs for RemoteInsertExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match t { @@ -278,6 +528,8 @@ impl ExecutionPlan for RemoteInsertExec { self.upload_id.clone(), self.tracker.clone(), self.branch.clone(), + self.max_bytes_per_request, + self.max_request_duration, ))) } @@ -310,8 +562,36 @@ impl ExecutionPlan for RemoteInsertExec { let upload_id = self.upload_id.clone(); let tracker = self.tracker.clone(); let branch = self.branch.clone(); + let max_bytes_per_request = self.max_bytes_per_request; + let max_request_duration = self.max_request_duration; let stream = futures::stream::once(async move { + // Multipart writes with a byte budget split the partition into + // several bounded, still-streamed requests so no single request + // stays open long enough to hit the client read timeout. + if let (Some(upload_id), Some(max_bytes)) = + (upload_id.as_deref(), max_bytes_per_request) + { + let ctx = PartRequestCtx { + client: &client, + identifier: &identifier, + table_name: &table_name, + upload_id, + branch: branch.as_deref(), + overwrite, + }; + ctx.send_multipart_chunked(max_bytes, max_request_duration, input_stream, tracker) + .await?; + // Count 0 here as for the non-multipart path below: the parts are + // only staged, so the real row count is resolved when the caller + // completes the multipart write. + let count_array: ArrayRef = Arc::new(UInt64Array::from(vec![0u64])); + return Ok::(RecordBatch::try_new( + COUNT_SCHEMA.clone(), + vec![count_array], + )?); + } + let mut request = client .post(&format!("/v1/table/{}/insert/", identifier)) .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE); @@ -423,9 +703,15 @@ mod tests { use arrow_schema::{DataType, Field, Schema as ArrowSchema}; use datafusion::prelude::SessionContext; use datafusion_catalog::MemTable; - use std::sync::Arc; + use datafusion_common::{DataFusionError, Result as DataFusionResult}; + use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + use datafusion_physical_expr::EquivalenceProperties; + use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + use super::RemoteInsertExec; use crate::Table; use crate::remote::ARROW_STREAM_CONTENT_TYPE; use crate::table::datafusion::BaseTableAdapter; @@ -591,4 +877,489 @@ mod tests { // Verify: should have made exactly one HTTP request despite multiple input partitions assert_eq!(request_count.load(Ordering::SeqCst), 1); } + + /// Build a single-partition input plan from the given batches. + async fn input_plan_from_batches( + schema: Arc, + batches: Vec, + ) -> Arc { + use datafusion_catalog::TableProvider; + let mem = MemTable::try_new(schema, vec![batches]).unwrap(); + let ctx = SessionContext::new(); + mem.scan(&ctx.state(), None, &[], None).await.unwrap() + } + + /// Build a single-partition input plan from the batches spread across the + /// given partitions. + async fn input_plan_from_partitions( + schema: Arc, + partitions: Vec>, + ) -> Arc { + use datafusion_catalog::TableProvider; + let mem = MemTable::try_new(schema, partitions).unwrap(); + let ctx = SessionContext::new(); + mem.scan(&ctx.state(), None, &[], None).await.unwrap() + } + + fn counting_insert_client( + counter: Arc, + ) -> crate::remote::client::RestfulLanceDbClient + { + crate::remote::client::test_utils::client_with_handler(move |request| { + let path = request.url().path(); + assert_eq!(path, "/v1/table/my_table/insert/"); + let query = request.url().query().unwrap_or(""); + assert!(query.contains("upload_id=upload-1"), "query: {query}"); + assert!(query.contains("upload_part_id="), "query: {query}"); + counter.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(200) + .body(String::new()) + .unwrap() + }) + } + + /// Insert handler that records the `upload_part_id` of every part request so + /// a test can assert the ids are distinct. + fn recording_insert_client( + part_ids: Arc>>, + ) -> crate::remote::client::RestfulLanceDbClient + { + crate::remote::client::test_utils::client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/insert/"); + let part_id = request + .url() + .query_pairs() + .find(|(k, _)| k == "upload_part_id") + .map(|(_, v)| v.into_owned()) + .expect("upload_part_id query param"); + part_ids.lock().unwrap().push(part_id); + http::Response::builder() + .status(200) + .body(String::new()) + .unwrap() + }) + } + + /// Single-partition input plan that yields one good batch and then an error, + /// for exercising the mid-part input-error abort path in `send_one_part`. + #[derive(Debug)] + struct ErroringExec { + schema: Arc, + properties: Arc, + } + + impl ErroringExec { + fn new() -> Self { + let schema = record_batch!(("id", Int32, [1, 2])).unwrap().schema(); + let properties = PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + datafusion_physical_plan::Partitioning::UnknownPartitioning(1), + datafusion_physical_plan::execution_plan::EmissionType::Incremental, + datafusion_physical_plan::execution_plan::Boundedness::Bounded, + ); + Self { + schema, + properties: Arc::new(properties), + } + } + } + + impl DisplayAs for ErroringExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter<'_>, + ) -> std::fmt::Result { + write!(f, "ErroringExec") + } + } + + impl ExecutionPlan for ErroringExec { + fn name(&self) -> &str { + "ErroringExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Ok(self) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> DataFusionResult { + let batch = record_batch!(("id", Int32, [1, 2])).unwrap(); + let stream = futures::stream::iter(vec![ + Ok(batch), + Err(DataFusionError::Execution("boom".to_string())), + ]); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) + } + } + + #[tokio::test] + async fn test_multipart_chunked_splits_into_parts() { + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + let batches = vec![ + record_batch!(("id", Int32, [1, 2])).unwrap(), + record_batch!(("id", Int32, [3, 4])).unwrap(), + record_batch!(("id", Int32, [5, 6])).unwrap(), + ]; + let input = input_plan_from_batches(schema, batches).await; + + // A 1-byte budget forces every batch into its own part. + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(1), + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + + assert_eq!(insert_count.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn test_multipart_single_part_when_under_budget() { + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + let batches = vec![ + record_batch!(("id", Int32, [1, 2])).unwrap(), + record_batch!(("id", Int32, [3, 4])).unwrap(), + record_batch!(("id", Int32, [5, 6])).unwrap(), + ]; + let input = input_plan_from_batches(schema, batches).await; + + // A large byte budget and no time limit keep the whole partition in a + // single part. + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(64 * 1024 * 1024), + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + + assert_eq!(insert_count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_multipart_chunked_splits_by_duration() { + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + let batches = vec![ + record_batch!(("id", Int32, [1, 2])).unwrap(), + record_batch!(("id", Int32, [3, 4])).unwrap(), + record_batch!(("id", Int32, [5, 6])).unwrap(), + ]; + let input = input_plan_from_batches(schema, batches).await; + + // A large byte budget but a tiny duration budget: writing and sending + // one batch already takes longer than the limit, so each batch is cut + // into its own part on the time check rather than the byte check. + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(64 * 1024 * 1024), + Some(std::time::Duration::from_nanos(1)), + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + + assert_eq!(insert_count.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn test_multipart_empty_partition_stages_nothing() { + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + // An empty partition should stage no parts; on the multipart path the + // write relies on another partition having data to commit. + let input = input_plan_from_batches(schema, vec![]).await; + + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(64 * 1024 * 1024), + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + + assert_eq!(insert_count.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn test_multipart_chunked_uses_distinct_part_ids() { + use futures::StreamExt; + use std::collections::HashSet; + + let part_ids = Arc::new(Mutex::new(Vec::new())); + let client = recording_insert_client(part_ids.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + let batches = vec![ + record_batch!(("id", Int32, [1, 2])).unwrap(), + record_batch!(("id", Int32, [3, 4])).unwrap(), + record_batch!(("id", Int32, [5, 6])).unwrap(), + ]; + let input = input_plan_from_batches(schema, batches).await; + + // A 1-byte budget forces every batch into its own part. + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(1), + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + + let ids = part_ids.lock().unwrap().clone(); + assert_eq!(ids.len(), 3, "expected one part id per part: {ids:?}"); + assert!( + ids.iter().all(|id| !id.is_empty()), + "part ids must be non-empty: {ids:?}" + ); + let unique: HashSet<&String> = ids.iter().collect(); + assert_eq!(unique.len(), 3, "part ids must be distinct: {ids:?}"); + } + + #[tokio::test] + async fn test_multipart_chunks_each_partition_independently() { + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + let partitions = vec![ + // Partition 0: two batches, split into two parts by the 1-byte budget. + vec![ + record_batch!(("id", Int32, [1, 2])).unwrap(), + record_batch!(("id", Int32, [3, 4])).unwrap(), + ], + // Partition 1: one batch, one part. + vec![record_batch!(("id", Int32, [5, 6])).unwrap()], + ]; + let input = input_plan_from_partitions(schema, partitions).await; + + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(1), + None, + ); + + for partition in 0..2 { + let mut stream = exec + .execute(partition, Arc::new(TaskContext::default())) + .unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + } + + // 2 parts from partition 0 + 1 part from partition 1. + assert_eq!(insert_count.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn test_multipart_input_error_surfaces_original() { + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + // A large byte budget keeps the good batch and the following error in + // the same part, exercising the mid-part abort path. + let input: Arc = Arc::new(ErroringExec::new()); + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + None, + None, + Some(64 * 1024 * 1024), + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + let mut err = None; + while let Some(item) = stream.next().await { + if let Err(e) = item { + err = Some(e); + break; + } + } + + let err = err.expect("expected the input stream error to surface"); + // The original DataFusion error must win over the HTTP error it induces. + assert!( + err.to_string().contains("boom"), + "expected original input error, got: {err}" + ); + } + + #[tokio::test] + async fn test_multipart_records_progress_within_a_part() { + use crate::table::write_progress::{ProgressCallback, WriteProgress, WriteProgressTracker}; + use futures::StreamExt; + + let insert_count = Arc::new(AtomicUsize::new(0)); + let client = counting_insert_client(insert_count.clone()); + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "id", + DataType::Int32, + true, + )])); + let batches = vec![ + record_batch!(("id", Int32, [1, 2])).unwrap(), + record_batch!(("id", Int32, [3, 4])).unwrap(), + record_batch!(("id", Int32, [5, 6])).unwrap(), + ]; + let input = input_plan_from_batches(schema, batches).await; + + let observed = Arc::new(Mutex::new(Vec::::new())); + let observed_cb = observed.clone(); + let callback: ProgressCallback = Arc::new(Mutex::new(move |p: &WriteProgress| { + observed_cb.lock().unwrap().push(p.output_bytes()); + })); + let tracker = Arc::new(WriteProgressTracker::new(callback, None)); + + // A large byte budget keeps all three batches in one part; smooth + // progress therefore requires bytes to be reported per chunk rather than + // once when the part completes. + let exec = RemoteInsertExec::new_multipart( + "my_table".to_string(), + "my_table".to_string(), + client, + input, + false, + "upload-1".to_string(), + Some(tracker), + None, + Some(64 * 1024 * 1024), + None, + ); + + let mut stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + while stream.next().await.transpose().unwrap().is_some() {} + + assert_eq!( + insert_count.load(Ordering::SeqCst), + 1, + "batches should all land in a single part" + ); + let observed = observed.lock().unwrap(); + assert!( + observed.len() > 1, + "expected multiple incremental progress updates within the part: {observed:?}" + ); + assert!( + observed.windows(2).all(|w| w[1] >= w[0]), + "progress bytes should be monotonic: {observed:?}" + ); + assert!( + *observed.last().unwrap() > 0, + "final progress should report bytes: {observed:?}" + ); + } }