diff --git a/Cargo.lock b/Cargo.lock index 5cd3b026197..c65a9675db7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1482,7 +1482,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2763,7 +2763,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2906,7 +2906,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4215,7 +4215,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi", "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4527,7 +4527,7 @@ dependencies = [ "lzma-sys", "mock_instant", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "parquet", "permutation", "pin-project", @@ -4643,7 +4643,7 @@ dependencies = [ "log", "moka", "num_cpus", - "object_store 0.14.1", + "object_store 0.14.2", "pin-project", "proptest", "prost", @@ -4788,7 +4788,7 @@ dependencies = [ "lance-datagen", "lance-index", "lance-linalg", - "object_store 0.14.1", + "object_store 0.14.2", "parquet", "rand 0.9.5", "tempfile", @@ -4823,7 +4823,7 @@ dependencies = [ "libc", "log", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "pretty_assertions", "proptest", "prost", @@ -4912,7 +4912,7 @@ dependencies = [ "log", "ndarray", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -4986,7 +4986,7 @@ dependencies = [ "mock_instant", "mockall", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "object_store_opendal", "opendal", "opendal-service-goosefs", @@ -5089,7 +5089,7 @@ dependencies = [ "lance-namespace", "lance-table", "log", - "object_store 0.14.1", + "object_store 0.14.2", "quick-xml 0.40.1", "rand 0.9.5", "reqwest 0.12.28", @@ -5165,7 +5165,7 @@ dependencies = [ "lance-select", "lance-testing", "log", - "object_store 0.14.1", + "object_store 0.14.2", "pretty_assertions", "proptest", "prost", @@ -5229,7 +5229,7 @@ dependencies = [ "lance-core", "lance-file", "lance-io", - "object_store 0.14.1", + "object_store 0.14.2", "tokio", "url", ] @@ -5874,7 +5874,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -6061,13 +6061,12 @@ dependencies = [ [[package]] name = "object_store" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" +version = "0.14.2" +source = "git+https://github.com/apache/arrow-rs-object-store.git?rev=c10758c207c0103355f940f0c736fb8ecb157b04#c10758c207c0103355f940f0c736fb8ecb157b04" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.0", "bytes", "chrono", "crc-fast", @@ -6113,7 +6112,7 @@ dependencies = [ "bytes", "chrono", "futures", - "object_store 0.14.1", + "object_store 0.14.2", "opendal", "pin-project", "tokio", @@ -7228,7 +7227,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.52.0", ] [[package]] @@ -7991,7 +7990,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -8050,7 +8049,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -8923,7 +8922,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -9519,7 +9518,7 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" dependencies = [ - "rand 0.10.1", + "rand 0.9.5", ] [[package]] @@ -9969,7 +9968,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index fc6f1aa552c..c4389620be5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -176,7 +176,7 @@ mock_instant = { version = "0.6.0" } moka = { version = "0.12", features = ["future", "sync"] } ndarray = { version = "0.16.1", features = ["matrixmultiply-threading"] } num-traits = "0.2" -object_store = { version = "0.14.1" } +object_store = { version = "0.14.2" } opendal = { version = "0.59.2" } object_store_opendal = { version = "0.60.2" } # Pin 0.2.2 so Cargo unifies past OpenDAL 0.59.2's `^0.2.1` req. 0.2.2 keeps @@ -283,3 +283,6 @@ multiple-crate-versions = "allow" single_range_in_vec_init = "allow" large_futures = "deny" disallowed_macros = "deny" + +[patch.crates-io] +object_store = { git = "https://github.com/apache/arrow-rs-object-store.git", rev = "c10758c207c0103355f940f0c736fb8ecb157b04" } diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 641fa84534c..a483d9fcab2 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -93,7 +93,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -104,7 +104,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1222,7 +1222,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2330,7 +2330,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2435,7 +2435,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3895,7 +3895,7 @@ dependencies = [ "lance-tokenizer", "log", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "permutation", "pin-project", "prost", @@ -3991,7 +3991,7 @@ dependencies = [ "log", "moka", "num_cpus", - "object_store 0.14.1", + "object_store 0.14.2", "pin-project", "prost", "quick_cache", @@ -4104,7 +4104,7 @@ dependencies = [ "lance-io", "log", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -4179,7 +4179,7 @@ dependencies = [ "log", "ndarray", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -4243,7 +4243,7 @@ dependencies = [ "log", "metrics", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "object_store_opendal", "opendal", "opendal-service-goosefs", @@ -4291,7 +4291,7 @@ dependencies = [ "log", "metrics", "metrics-util", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-types", "roaring", @@ -4351,7 +4351,7 @@ dependencies = [ "lance-namespace", "lance-table", "log", - "object_store 0.14.1", + "object_store 0.14.2", "rand 0.9.5", "reqwest 0.12.28", "roaring", @@ -4413,7 +4413,7 @@ dependencies = [ "lance-io", "lance-select", "log", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -4934,7 +4934,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -5096,13 +5096,12 @@ dependencies = [ [[package]] name = "object_store" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" +version = "0.14.2" +source = "git+https://github.com/apache/arrow-rs-object-store.git?rev=c10758c207c0103355f940f0c736fb8ecb157b04#c10758c207c0103355f940f0c736fb8ecb157b04" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.0", "bytes", "chrono", "crc-fast", @@ -5148,7 +5147,7 @@ dependencies = [ "bytes", "chrono", "futures", - "object_store 0.14.1", + "object_store 0.14.2", "opendal", "pin-project", "tokio", @@ -6657,7 +6656,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -6715,7 +6714,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7203,7 +7202,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -7468,7 +7467,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7932,7 +7931,7 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" dependencies = [ - "rand 0.10.1", + "rand 0.9.5", ] [[package]] @@ -8308,7 +8307,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/java/lance-jni/Cargo.toml b/java/lance-jni/Cargo.toml index 68631cffbcc..685e059cee7 100644 --- a/java/lance-jni/Cargo.toml +++ b/java/lance-jni/Cargo.toml @@ -39,7 +39,7 @@ arrow-array = "58.0.0" arrow-schema = "58.0.0" datafusion = { version = "54.0.0", default-features = false } datafusion-common = "54.0.0" -object_store = { version = "0.14.1" } +object_store = { version = "0.14.2" } tokio = { version = "1.23", features = [ "rt-multi-thread", "macros", @@ -79,3 +79,6 @@ strip = true [lints.clippy] disallowed_macros = "deny" + +[patch.crates-io] +object_store = { git = "https://github.com/apache/arrow-rs-object-store.git", rev = "c10758c207c0103355f940f0c736fb8ecb157b04" } diff --git a/python/Cargo.lock b/python/Cargo.lock index 3a420661f2a..48864c9124f 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -108,7 +108,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -119,7 +119,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1298,7 +1298,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -2509,7 +2509,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2620,7 +2620,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4063,7 +4063,7 @@ dependencies = [ "lance-tokenizer", "log", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "permutation", "pin-project", "prost", @@ -4159,7 +4159,7 @@ dependencies = [ "log", "moka", "num_cpus", - "object_store 0.14.1", + "object_store 0.14.2", "pin-project", "prost", "quick_cache", @@ -4289,7 +4289,7 @@ dependencies = [ "lance-io", "log", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -4364,7 +4364,7 @@ dependencies = [ "log", "ndarray", "num-traits", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -4428,7 +4428,7 @@ dependencies = [ "log", "metrics", "moka", - "object_store 0.14.1", + "object_store 0.14.2", "object_store_opendal", "opendal", "opendal-service-goosefs", @@ -4498,7 +4498,7 @@ dependencies = [ "lance-namespace", "lance-table", "log", - "object_store 0.14.1", + "object_store 0.14.2", "rand 0.9.5", "reqwest 0.12.28", "roaring", @@ -4562,7 +4562,7 @@ dependencies = [ "lance-io", "lance-select", "log", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-build", "prost-types", @@ -5107,7 +5107,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -5269,13 +5269,12 @@ dependencies = [ [[package]] name = "object_store" -version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" +version = "0.14.2" +source = "git+https://github.com/apache/arrow-rs-object-store.git?rev=c10758c207c0103355f940f0c736fb8ecb157b04#c10758c207c0103355f940f0c736fb8ecb157b04" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.0", "bytes", "chrono", "crc-fast", @@ -5321,7 +5320,7 @@ dependencies = [ "bytes", "chrono", "futures", - "object_store 0.14.1", + "object_store 0.14.2", "opendal", "pin-project", "tokio", @@ -6175,7 +6174,7 @@ dependencies = [ "log", "metrics", "metrics-util", - "object_store 0.14.1", + "object_store 0.14.2", "prost", "prost-types", "pyo3", @@ -7040,7 +7039,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7098,7 +7097,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7598,7 +7597,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -7903,7 +7902,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -8389,7 +8388,7 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" dependencies = [ - "rand 0.10.1", + "rand 0.9.5", ] [[package]] @@ -8765,7 +8764,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/python/Cargo.toml b/python/Cargo.toml index cabe5a0c180..2ff9b1da84a 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -24,7 +24,7 @@ arrow-array = "58.0.0" arrow-cast = "58.0.0" arrow-data = "58.0.0" arrow-schema = "58.0.0" -object_store = "0.14.1" +object_store = "0.14.2" url = "2.5.7" datafusion = { version = "54.0.0", default-features = false } datafusion-ffi = "54.0.0" @@ -105,5 +105,8 @@ debug-assertions = false strip = "debuginfo" incremental = false +[patch.crates-io] +object_store = { git = "https://github.com/apache/arrow-rs-object-store.git", rev = "c10758c207c0103355f940f0c736fb8ecb157b04" } + [lints.clippy] disallowed_macros = "deny" diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 21cf127a321..fdf00e94343 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -84,6 +84,12 @@ pub const DEFAULT_CLOUD_IO_PARALLELISM: usize = 64; const SERVER_SIDE_COPY_ENABLED_ENV: &str = "LANCE_IO_SERVER_SIDE_COPY_ENABLED"; +#[derive(Clone, Copy)] +enum StreamCopyMode { + Overwrite, + Create, +} + const DEFAULT_LOCAL_BLOCK_SIZE: usize = 4 * 1024; // 4KB block size #[cfg(any( feature = "aws", @@ -211,6 +217,8 @@ pub struct ObjectStore { /// Whether to use constant size upload parts for multipart uploads. This /// is only necessary for Cloudflare R2. pub use_constant_size_upload_parts: bool, + /// Whether multipart uploads honor [`PutMode::Create`] at completion. + supports_conditional_multipart_put: bool, /// Whether we can assume that the list of files is lexically ordered. This /// is true for object stores, but not for local filesystems. pub list_is_lexically_ordered: bool, @@ -240,6 +248,10 @@ impl std::fmt::Debug for ObjectStore { "use_constant_size_upload_parts", &self.use_constant_size_upload_parts, ) + .field( + "supports_conditional_multipart_put", + &self.supports_conditional_multipart_put, + ) .field("list_is_lexically_ordered", &self.list_is_lexically_ordered) .field("io_parallelism", &self.io_parallelism) .field("download_retry_count", &self.download_retry_count) @@ -688,6 +700,7 @@ impl ObjectStore { 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, + supports_conditional_multipart_put: false, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or_default(), io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT, @@ -850,6 +863,7 @@ impl ObjectStore { /// [`Self::read_dir_page`] falls back through the (already-wrapped) `inner`. pub fn apply_wrapper(&mut self, wrapper: &dyn WrappingObjectStore) { self.inner = wrapper.wrap(&self.store_prefix, self.inner.clone()); + self.supports_conditional_multipart_put = false; self.paginated_lister = self .paginated_lister .take() @@ -1247,6 +1261,103 @@ impl ObjectStore { source_path: &Path, destination_store: &Self, destination_path: &Path, + ) -> Result { + self.copy_via_stream_with_mode( + source_path, + destination_store, + destination_path, + StreamCopyMode::Overwrite, + ) + .await + } + + /// Copy within this store without replacing an existing destination. + /// + /// Backends with verified conditional multipart completion stream through + /// the client. Other backends retain their native conditional-copy path. + pub async fn copy_if_not_exists( + &self, + source_path: &Path, + destination_path: &Path, + ) -> Result { + if self.supports_conditional_multipart_put { + return self + .copy_if_not_exists_via_stream(source_path, self, destination_path) + .await; + } + + let source_size = self.size(source_path).await?; + match self + .inner + .copy_if_not_exists(source_path, destination_path) + .await + { + Ok(()) => {} + Err( + object_store::Error::AlreadyExists { .. } + | object_store::Error::Precondition { .. }, + ) => { + let destination_size = self.size(destination_path).await?; + if source_size != destination_size { + return Err(Error::io(format!( + "destination {destination_path} already exists with size \ + {destination_size}, but source {source_path} has size {source_size}" + ))); + } + } + Err(error) => return Err(Error::io(error.to_string())), + } + + Ok(WriteResult { + size: usize::try_from(source_size).map_err(|_| { + Error::io(format!( + "source {source_path} size {source_size} does not fit in usize" + )) + })?, + e_tag: None, + }) + } + + #[instrument( + name = "multipart_stream_copy_if_not_exists", + level = "info", + skip(self, source_path, destination_store, destination_path), + fields( + source = %source_path, + destination = %destination_path, + source_size = Empty, + read_chunk_size = Empty, + multipart_part_size = crate::object_writer::initial_upload_size(), + multipart_concurrency = crate::object_writer::max_upload_parallelism(), + part_count = Empty, + bytes_transferred = Empty, + destination_size = Empty, + validation = Empty, + elapsed_ms = Empty, + ), + err + )] + async fn copy_if_not_exists_via_stream( + &self, + source_path: &Path, + destination_store: &Self, + destination_path: &Path, + ) -> Result { + self.copy_via_stream_with_mode( + source_path, + destination_store, + destination_path, + StreamCopyMode::Create, + ) + .await + } + + async fn copy_via_stream_with_mode( + &self, + source_path: &Path, + destination_store: &Self, + destination_path: &Path, + mode: StreamCopyMode, ) -> Result { let started_at = Instant::now(); if self.has_direct_local_paths() && destination_store.has_direct_local_paths() { @@ -1271,7 +1382,20 @@ impl ObjectStore { Span::current().record("source_size", source_size as u64); let metrics = destination_store.io_tracker.begin_io("copy"); - let result = super::local::copy_file(source_path, destination_path); + let result = match mode { + StreamCopyMode::Overwrite => super::local::copy_file(source_path, destination_path), + StreamCopyMode::Create => match destination_store + .inner + .copy_if_not_exists(source_path, destination_path) + .await + { + Ok(()) => Ok(()), + Err(copy_error) => match destination_store.size(destination_path).await { + Ok(destination_size) if destination_size == source_size as u64 => Ok(()), + _ => Err(Error::from(copy_error)), + }, + }, + }; metrics.record(&result, source_size as u64); result.map_err(|source| { stream_copy_error( @@ -1321,17 +1445,22 @@ impl ObjectStore { })?; Span::current().record("source_size", source_size as u64); - let mut writer = destination_store - .create(destination_path) - .await - .map_err(|source| { - stream_copy_error( - "destination writer creation", - source_path, - destination_path, - source, - ) - })?; + let mut writer: Box = match mode { + StreamCopyMode::Overwrite => destination_store.create(destination_path).await, + StreamCopyMode::Create => { + ObjectWriter::new_with_mode(destination_store, destination_path, PutMode::Create) + .await + .map(|writer| Box::new(writer) as Box) + } + } + .map_err(|source| { + stream_copy_error( + "destination writer creation", + source_path, + destination_path, + source, + ) + })?; let read_chunk_size = usize::try_from(self.max_iop_size()) .unwrap_or(usize::MAX) .max(1); @@ -1395,14 +1524,36 @@ impl ObjectStore { } Span::current().record("bytes_transferred", bytes_transferred as u64); - let write_result = Writer::shutdown(writer.as_mut()).await.map_err(|source| { - stream_copy_error( - "destination completion", - source_path, - destination_path, - source, - ) - })?; + let write_result = match Writer::shutdown(writer.as_mut()).await { + Ok(write_result) => write_result, + Err(source) => { + let completion_error = stream_copy_error( + "destination completion", + source_path, + destination_path, + source, + ); + if !matches!(mode, StreamCopyMode::Create) { + return Err(completion_error); + } + let Ok(destination_size) = destination_store.size(destination_path).await else { + return Err(completion_error); + }; + Span::current().record("destination_size", destination_size); + if destination_size != source_size as u64 { + Span::current().record("validation", "failed"); + return Err(Error::io(format!( + "multipart_stream_copy existing destination size mismatch from \ + {source_path} to {destination_path}: source_size={source_size}, \ + destination_size={destination_size}" + ))); + } + WriteResult { + size: source_size, + e_tag: None, + } + } + }; if write_result.size != source_size { Span::current().record("validation", "failed"); return Err(Error::io(format!( @@ -1864,6 +2015,7 @@ impl ObjectStore { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts, + supports_conditional_multipart_put: false, list_is_lexically_ordered, io_parallelism, download_retry_count, @@ -3032,6 +3184,136 @@ mod tests { ); } + #[tokio::test] + async fn test_copy_if_not_exists_via_stream_preserves_existing_destination() { + let mut store = ObjectStore::memory(); + store.inner = Arc::new(CopyFailingStore { + inner: InMemory::new(), + }); + + let source = Path::from("source.bin"); + let destination = Path::from("destination.bin"); + store.put(&source, b"new contents").await.unwrap(); + store.put(&destination, b"existing contents").await.unwrap(); + + store + .copy_if_not_exists_via_stream(&source, &store, &destination) + .await + .expect_err("conditional stream copy must not replace an existing object"); + + assert_eq!( + store.read_one_all(&destination).await.unwrap().as_ref(), + b"existing contents" + ); + assert_eq!( + store.read_one_all(&source).await.unwrap().as_ref(), + b"new contents" + ); + } + + #[tokio::test] + async fn test_copy_if_not_exists_via_stream_creates_destination() { + let mut store = ObjectStore::memory(); + store.inner = Arc::new(CopyFailingStore { + inner: InMemory::new(), + }); + + let source = Path::from("source.bin"); + let destination = Path::from("destination.bin"); + let contents = b"stream into a previously absent destination"; + store.put(&source, contents).await.unwrap(); + + let result = store + .copy_if_not_exists_via_stream(&source, &store, &destination) + .await + .unwrap(); + + assert_eq!(result.size, contents.len()); + assert_eq!( + store.read_one_all(&destination).await.unwrap().as_ref(), + contents + ); + } + + #[tokio::test] + async fn test_copy_if_not_exists_via_stream_accepts_equal_size_retry() { + let mut store = ObjectStore::memory(); + store.inner = Arc::new(CopyFailingStore { + inner: InMemory::new(), + }); + + let source = Path::from("source.bin"); + let destination = Path::from("destination.bin"); + store.put(&source, b"new").await.unwrap(); + store.put(&destination, b"old").await.unwrap(); + + let result = store + .copy_if_not_exists_via_stream(&source, &store, &destination) + .await + .unwrap(); + + assert_eq!(result.size, 3); + assert_eq!( + store.read_one_all(&destination).await.unwrap().as_ref(), + b"old" + ); + } + + #[tokio::test] + async fn test_copy_if_not_exists_accepts_equal_size_local_retry() { + let directory = TempStdDir::default(); + let (store, base_path) = ObjectStore::from_uri(directory.to_str().unwrap()) + .await + .unwrap(); + let source = base_path.clone().join("source.bin"); + let destination = base_path.join("destination.bin"); + store.put(&source, b"new").await.unwrap(); + store.put(&destination, b"old").await.unwrap(); + + let result = store + .copy_if_not_exists(&source, &destination) + .await + .unwrap(); + + assert_eq!(result.size, 3); + assert_eq!( + store.read_one_all(&destination).await.unwrap().as_ref(), + b"old" + ); + } + + #[tokio::test] + async fn test_copy_if_not_exists_uses_native_copy_without_conditional_multipart() { + let observations = Arc::new(MultipartObservations::default()); + let mut store = ObjectStore::memory(); + store.supports_conditional_multipart_put = false; + store.inner = Arc::new(ObservedMultipartStore { + inner: InMemory::new(), + observations: observations.clone(), + fail_parts: false, + destination_size_adjustment: 0, + }); + + let source = Path::from("source.bin"); + let destination = Path::from("destination.bin"); + let contents = vec![42; crate::object_writer::initial_upload_size() + 1]; + store.put(&source, &contents).await.unwrap(); + observations.part_count.store(0, Ordering::SeqCst); + + let result = store + .copy_if_not_exists(&source, &destination) + .await + .unwrap(); + + assert_eq!(result.size, contents.len()); + assert_eq!(observations.native_copy_count.load(Ordering::SeqCst), 1); + assert_eq!(observations.part_count.load(Ordering::SeqCst), 0); + assert_eq!( + store.read_one_all(&destination).await.unwrap().as_ref(), + contents.as_slice() + ); + } + #[tokio::test] async fn test_remove_stream_removes_every_path() { let store = ObjectStore::memory(); @@ -3382,6 +3664,52 @@ mod tests { ); } + #[tokio::test] + async fn test_copy_if_not_exists_via_stream_uses_conditional_multipart_completion() { + let mut source_store = ObjectStore::memory(); + source_store.max_iop_size = 1024 * 1024; + let observations = Arc::new(MultipartObservations::default()); + let mut destination_store = ObjectStore::memory(); + destination_store.inner = Arc::new(ObservedMultipartStore { + inner: InMemory::new(), + observations: observations.clone(), + fail_parts: false, + destination_size_adjustment: 0, + }); + + let source = Path::from("source.bin"); + let destination = Path::from("destination.bin"); + let contents = vec![42; crate::object_writer::initial_upload_size() * 2 + 1]; + source_store.put(&source, &contents).await.unwrap(); + destination_store + .put(&destination, b"existing contents") + .await + .unwrap(); + + source_store + .copy_if_not_exists_via_stream(&source, &destination_store, &destination) + .await + .expect_err("conditional multipart completion must reject an existing object"); + + assert!( + observations.part_count.load(Ordering::SeqCst) >= 2, + "conditional stream copy should use multipart upload for a large destination" + ); + assert_eq!( + observations.abort_count.load(Ordering::SeqCst), + 1, + "conditional completion failure should abort the multipart upload" + ); + assert_eq!( + destination_store + .read_one_all(&destination) + .await + .unwrap() + .as_ref(), + b"existing contents" + ); + } + #[tokio::test] async fn test_copy_via_stream_aborts_failed_upload_and_retains_source() { let source_store = ObjectStore::memory(); diff --git a/rust/lance-io/src/object_store/providers.rs b/rust/lance-io/src/object_store/providers.rs index 4982c2b97ee..12bdc868a67 100644 --- a/rust/lance-io/src/object_store/providers.rs +++ b/rust/lance-io/src/object_store/providers.rs @@ -551,6 +551,7 @@ mod tests { .unwrap(); assert_eq!(store.paginated_lister.is_some(), keep_pushdown); + assert!(!store.supports_conditional_multipart_put); // Both halves of the store are labelled with the same prefix. assert_eq!( *wrapper.prefixes.lock().unwrap(), diff --git a/rust/lance-io/src/object_store/providers/aws.rs b/rust/lance-io/src/object_store/providers/aws.rs index 3702e6be5ae..b638a914c32 100644 --- a/rust/lance-io/src/object_store/providers/aws.rs +++ b/rust/lance-io/src/object_store/providers/aws.rs @@ -258,6 +258,7 @@ impl ObjectStoreProvider for AwsStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts, + supports_conditional_multipart_put: !use_opendal, list_is_lexically_ordered: !is_s3_express, io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count, @@ -1073,6 +1074,7 @@ mod tests { .await .unwrap(); assert_eq!(store.scheme, "s3"); + assert!(!store.supports_conditional_multipart_put); } /// S3 Express ignores `start-after` and does not list in key order, but it does hand back @@ -1103,6 +1105,10 @@ mod tests { assert!(!store.list_is_lexically_ordered); assert_eq!(store.paginated_lister.is_some(), paginated); + assert_eq!( + store.supports_conditional_multipart_put, + use_opendal == "false" + ); } #[derive(Debug)] diff --git a/rust/lance-io/src/object_store/providers/azure.rs b/rust/lance-io/src/object_store/providers/azure.rs index 794a717b31c..a86e7e2ff27 100644 --- a/rust/lance-io/src/object_store/providers/azure.rs +++ b/rust/lance-io/src/object_store/providers/azure.rs @@ -300,6 +300,7 @@ impl ObjectStoreProvider for AzureBlobStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: false, + supports_conditional_multipart_put: false, list_is_lexically_ordered: true, io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count, diff --git a/rust/lance-io/src/object_store/providers/gcp.rs b/rust/lance-io/src/object_store/providers/gcp.rs index e3e4e4ee81c..4874d4b022c 100644 --- a/rust/lance-io/src/object_store/providers/gcp.rs +++ b/rust/lance-io/src/object_store/providers/gcp.rs @@ -344,6 +344,7 @@ impl ObjectStoreProvider for GcsStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: false, + supports_conditional_multipart_put: false, list_is_lexically_ordered: true, io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count, diff --git a/rust/lance-io/src/object_store/providers/goosefs.rs b/rust/lance-io/src/object_store/providers/goosefs.rs index ea13c7d38b6..59217432a30 100644 --- a/rust/lance-io/src/object_store/providers/goosefs.rs +++ b/rust/lance-io/src/object_store/providers/goosefs.rs @@ -489,6 +489,7 @@ impl ObjectStoreProvider for GooseFsStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: params.use_constant_size_upload_parts, + supports_conditional_multipart_put: false, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(false), io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count: storage_options.download_retry_count(), diff --git a/rust/lance-io/src/object_store/providers/huggingface.rs b/rust/lance-io/src/object_store/providers/huggingface.rs index dac76d130ef..5e6caf84cf8 100644 --- a/rust/lance-io/src/object_store/providers/huggingface.rs +++ b/rust/lance-io/src/object_store/providers/huggingface.rs @@ -226,6 +226,7 @@ impl ObjectStoreProvider for HuggingfaceStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: params.use_constant_size_upload_parts, + supports_conditional_multipart_put: false, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count, diff --git a/rust/lance-io/src/object_store/providers/local.rs b/rust/lance-io/src/object_store/providers/local.rs index 4dd834c9151..03b6437fdf9 100644 --- a/rust/lance-io/src/object_store/providers/local.rs +++ b/rust/lance-io/src/object_store/providers/local.rs @@ -132,6 +132,7 @@ impl ObjectStoreProvider for FileStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: false, + supports_conditional_multipart_put: false, list_is_lexically_ordered: false, io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM, download_retry_count, @@ -198,6 +199,7 @@ mod tests { block_size: DEFAULT_LOCAL_BLOCK_SIZE, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: false, + supports_conditional_multipart_put: false, list_is_lexically_ordered: false, io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM, download_retry_count: 0, diff --git a/rust/lance-io/src/object_store/providers/memory.rs b/rust/lance-io/src/object_store/providers/memory.rs index f9b4a22e2cc..7d9f33ef565 100644 --- a/rust/lance-io/src/object_store/providers/memory.rs +++ b/rust/lance-io/src/object_store/providers/memory.rs @@ -28,6 +28,7 @@ impl ObjectStoreProvider for MemoryStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: false, + supports_conditional_multipart_put: true, list_is_lexically_ordered: true, io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count, diff --git a/rust/lance-io/src/object_store/providers/oss.rs b/rust/lance-io/src/object_store/providers/oss.rs index 0adf52db0a2..d2d07792a4b 100644 --- a/rust/lance-io/src/object_store/providers/oss.rs +++ b/rust/lance-io/src/object_store/providers/oss.rs @@ -140,6 +140,7 @@ impl ObjectStoreProvider for OssStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: params.use_constant_size_upload_parts, + supports_conditional_multipart_put: false, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count: storage_options.download_retry_count(), diff --git a/rust/lance-io/src/object_store/providers/tencent.rs b/rust/lance-io/src/object_store/providers/tencent.rs index 9ad0a91765d..144df51793c 100644 --- a/rust/lance-io/src/object_store/providers/tencent.rs +++ b/rust/lance-io/src/object_store/providers/tencent.rs @@ -96,6 +96,7 @@ impl ObjectStoreProvider for TencentStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: params.use_constant_size_upload_parts, + supports_conditional_multipart_put: false, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count: storage_options.download_retry_count(), diff --git a/rust/lance-io/src/object_store/providers/tos.rs b/rust/lance-io/src/object_store/providers/tos.rs index 9f558e92550..ad296e0596f 100644 --- a/rust/lance-io/src/object_store/providers/tos.rs +++ b/rust/lance-io/src/object_store/providers/tos.rs @@ -146,6 +146,7 @@ impl ObjectStoreProvider for TosStoreProvider { block_size, max_iop_size: *DEFAULT_MAX_IOP_SIZE, use_constant_size_upload_parts: params.use_constant_size_upload_parts, + supports_conditional_multipart_put: false, list_is_lexically_ordered: params.list_is_lexically_ordered.unwrap_or(true), io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM, download_retry_count: storage_options.download_retry_count(), diff --git a/rust/lance-io/src/object_writer.rs b/rust/lance-io/src/object_writer.rs index c5eaffde99a..605011cb2be 100644 --- a/rust/lance-io/src/object_writer.rs +++ b/rust/lance-io/src/object_writer.rs @@ -12,7 +12,7 @@ use async_trait::async_trait; use bytes::Bytes; use futures::FutureExt; use futures::future::BoxFuture; -use object_store::{MultipartUpload, ObjectStoreExt}; +use object_store::{MultipartUpload, PutMode, PutMultipartOptions, PutOptions}; use object_store::{ObjectStore, path::Path}; use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio::task::JoinSet; @@ -90,6 +90,7 @@ pub struct ObjectWriter { path: Arc, cursor: usize, buffer: Vec, + put_mode: PutMode, // TODO: use constant size to support R2 use_constant_size_upload_parts: bool, } @@ -223,7 +224,7 @@ enum UploadState { /// Methods for state transitions. impl UploadState { - fn started_to_putting_single(&mut self, path: Arc, buffer: Vec) { + fn started_to_putting_single(&mut self, path: Arc, buffer: Vec, put_mode: PutMode) { // To get owned self, we temporarily swap with Done. let this = std::mem::replace(self, Self::Done(WriteResult::default())); *self = match this { @@ -232,16 +233,26 @@ impl UploadState { let started_at = Instant::now(); let fut = async move { let size = buffer.len(); - let res = store.put(&path, buffer.into()).await.map_err(|source| { - UploadFailure::new( - format!( - "single PUT of {path} failed after {:?} ({size} bytes, {})", - started_at.elapsed(), - upload_settings() - ), - source, + let res = store + .put_opts( + &path, + buffer.into(), + PutOptions { + mode: put_mode, + ..Default::default() + }, ) - })?; + .await + .map_err(|source| { + UploadFailure::new( + format!( + "single PUT of {path} failed after {:?} ({size} bytes, {})", + started_at.elapsed(), + upload_settings() + ), + source, + ) + })?; Ok(WriteResult { size, e_tag: res.e_tag, @@ -266,17 +277,28 @@ impl UploadState { tracing::Span::current().record("part_count", part_idx as u64); let started_at = Instant::now(); let fut = async move { - let res = upload.complete().await.map_err(|source| { - UploadFailure::new( - format!( - "completing multipart upload of {path} failed after {:?} \ + let res = match upload.complete().await { + Ok(result) => result, + Err(source) => { + let elapsed = started_at.elapsed(); + if let Err(error) = upload.abort().await { + tracing::warn!( + path = %path, + error = %error, + "Failed to abort multipart upload after completion error" + ); + } + return Err(UploadFailure::new( + format!( + "completing multipart upload of {path} failed after {:?} \ ({part_idx} parts, {bytes_written} bytes, {})", - started_at.elapsed(), - upload_settings() - ), - source, - ) - })?; + elapsed, + upload_settings() + ), + source, + )); + } + }; Ok(WriteResult { size: 0, // This will be set properly later. e_tag: res.e_tag, @@ -291,11 +313,20 @@ impl UploadState { impl ObjectWriter { pub async fn new(object_store: &LanceObjectStore, path: &Path) -> Result { + Self::new_with_mode(object_store, path, PutMode::Overwrite).await + } + + pub(crate) async fn new_with_mode( + object_store: &LanceObjectStore, + path: &Path, + put_mode: PutMode, + ) -> Result { Ok(Self { state: UploadState::Started(object_store.inner.clone()), cursor: 0, path: Arc::new(path.clone()), buffer: Vec::with_capacity(initial_upload_size()), + put_mode, use_constant_size_upload_parts: object_store.use_constant_size_upload_parts, }) } @@ -472,18 +503,28 @@ impl AsyncWrite for ObjectWriter { UploadState::Started(store) => { let path = mut_self.path.clone(); let store = store.clone(); + let put_mode = mut_self.put_mode.clone(); let started_at = Instant::now(); let fut = Box::pin(async move { - store.put_multipart(path.as_ref()).await.map_err(|source| { - UploadFailure::new( - format!( - "failed to create multipart upload for {path} after {:?} ({})", - started_at.elapsed(), - upload_settings() - ), - source, + store + .put_multipart_opts( + path.as_ref(), + PutMultipartOptions { + mode: put_mode, + ..Default::default() + }, ) - }) + .await + .map_err(|source| { + UploadFailure::new( + format!( + "failed to create multipart upload for {path} after {:?} ({})", + started_at.elapsed(), + upload_settings() + ), + source, + ) + }) }); self.state = UploadState::CreatingUpload(fut); } @@ -570,7 +611,8 @@ impl AsyncWrite for ObjectWriter { // If we didn't start a multipart upload, we can just do a single put. let part = std::mem::take(&mut mut_self.buffer); let path = mut_self.path.clone(); - self.state.started_to_putting_single(path, part); + let put_mode = mut_self.put_mode.clone(); + self.state.started_to_putting_single(path, part, put_mode); } UploadState::InProgress { upload, diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 15c49ae0c45..16a20932b2c 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -2100,38 +2100,9 @@ pub(in crate::dataset) struct FilePromotion { impl FilePromotion { async fn promote(&self) -> Result<()> { - match self - .object_store - .inner + self.object_store .copy_if_not_exists(&self.staging_path, &self.final_path) - .await - { - Ok(()) => {} - Err( - object_store::Error::AlreadyExists { .. } - | object_store::Error::Precondition { .. }, - ) => { - let staged = self - .object_store - .inner - .head(&self.staging_path) - .await - .map_err(|error| Error::io(error.to_string()))?; - let existing = self - .object_store - .inner - .head(&self.final_path) - .await - .map_err(|error| Error::io(error.to_string()))?; - if staged.size != existing.size { - return Err(Error::io(format!( - "fixed data file {} already exists with size {}, but the staged retry has size {}", - self.final_path, existing.size, staged.size - ))); - } - } - Err(error) => return Err(Error::io(error.to_string())), - } + .await?; self.object_store .inner .delete(&self.staging_path)