diff --git a/docs/src/guide/object_store.md b/docs/src/guide/object_store.md index cdc77af9a76..ace61d37338 100644 --- a/docs/src/guide/object_store.md +++ b/docs/src/guide/object_store.md @@ -30,6 +30,7 @@ These options apply to all object stores. | Key | Description | |------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| | `allow_http` | Allow non-TLS, i.e. non-HTTPS connections. Default, `False`. | +| `block_size` | Gap in bytes below which two reads of one file are merged into a single request. Default, `4096` for local files and `65536` for object stores. Larger values cut the request count of scattered reads (for example a `take` of rows spread across a file) at the cost of reading the bytes in between. | | `download_retry_count` | Number of times to retry a download. Default, `3`. This limit is applied when the HTTP request succeeds but the response is not fully downloaded, typically due to a violation of `timeout`. | | `allow_invalid_certificates` | Skip certificate validation on https connections. Default, `False`. Warning: This is insecure and should only be used for testing. | | `connect_timeout` | Timeout for only the connect phase of a Client. Default, `5s`. | diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 53a9767fb66..fc52c451b50 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -396,6 +396,18 @@ impl ObjectStoreParams { .and_then(|a| a.initial_storage_options()) } + /// The block size to use: the explicit `block_size` parameter, else the + /// `block_size` storage option, else `None` for the store's default. + pub fn resolved_block_size(&self) -> Result> { + if self.block_size.is_some() { + return Ok(self.block_size); + } + match self.storage_options() { + Some(options) => StorageOptions(options.clone()).block_size(), + None => Ok(None), + } + } + /// Resolve these params for a single base path scope. /// /// Storage options may carry base-scoped entries (`base_.`) that @@ -669,7 +681,7 @@ impl ObjectStore { inner: tracked_store, local_dir_operations: None, scheme: path.scheme().to_string(), - block_size: params.block_size.unwrap_or(64 * 1024), + block_size: params.resolved_block_size()?.unwrap_or(64 * 1024), max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: params.use_constant_size_upload_parts, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or_default(), @@ -1683,6 +1695,24 @@ impl StorageOptions { }) } + /// Byte gap below which the I/O scheduler merges two reads of one file + /// into a single request, overriding the store's default. + pub fn block_size(&self) -> Result> { + let Some((_, value)) = self + .0 + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("block_size")) + else { + return Ok(None); + }; + let block_size = value.trim().parse::().map_err(|err| { + Error::invalid_input(format!( + "storage option block_size must be a number of bytes, got `{value}`: {err}" + )) + })?; + Ok(Some(block_size)) + } + /// Number of times to retry a download that fails pub fn download_retry_count(&self) -> usize { self.0 @@ -2057,6 +2087,55 @@ mod tests { .await .unwrap(); assert_eq!(store.block_size, 1024); + + // The storage option applies when no parameter is given... + let mut options_with_block_size = storage_options.unwrap_or_default(); + options_with_block_size.insert(String::from("block_size"), String::from("2048")); + let accessor = Arc::new(StorageOptionsAccessor::with_static_options( + options_with_block_size, + )); + let registry = Arc::new(ObjectStoreRegistry::default()); + let params = ObjectStoreParams { + storage_options_accessor: Some(accessor.clone()), + ..ObjectStoreParams::default() + }; + let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms) + .await + .unwrap(); + assert_eq!(store.block_size, 2048); + + // ...and the explicit parameter wins over it. + let registry = Arc::new(ObjectStoreRegistry::default()); + let params = ObjectStoreParams { + block_size: Some(1024), + storage_options_accessor: Some(accessor), + ..ObjectStoreParams::default() + }; + let (store, _) = ObjectStore::from_uri_and_params(registry, uri, ¶ms) + .await + .unwrap(); + assert_eq!(store.block_size, 1024); + } + + #[tokio::test] + async fn test_block_size_option_rejects_invalid_values() { + let registry = Arc::new(ObjectStoreRegistry::default()); + let accessor = Arc::new(StorageOptionsAccessor::with_static_options(HashMap::from( + [(String::from("block_size"), String::from("64KiB"))], + ))); + let params = ObjectStoreParams { + storage_options_accessor: Some(accessor), + ..ObjectStoreParams::default() + }; + let error = + ObjectStore::from_uri_and_params(registry, "memory:///bucket/foo.lance", ¶ms) + .await + .unwrap_err(); + assert!( + matches!(error, lance_core::Error::InvalidInput { .. }), + "{error:?}" + ); + assert!(error.to_string().contains("block_size"), "{error}"); } #[rstest] diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index d8d184e07da..e29ec4a8d28 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -213,6 +213,11 @@ impl ObjectStoreRegistry { store_prefix: &str, ) -> Result> { let mut store = provider.new_store(base_path, params).await?; + // Providers only know the explicit parameter; the storage option is + // applied here so every store honours it the same way. + if let Some(block_size) = params.resolved_block_size()? { + store.block_size = block_size; + } store.inner = store.inner.traced();