diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs index 389f1a355ee..23c61f92dc0 100644 --- a/rust/lance-io/src/object_store.rs +++ b/rust/lance-io/src/object_store.rs @@ -45,6 +45,17 @@ pub(crate) mod dynamic_opendal; mod list_retry; #[cfg(feature = "metrics")] pub mod metrics; +#[cfg(any( + feature = "aws", + feature = "gcp", + feature = "azure", + feature = "oss", + feature = "tencent", + feature = "huggingface", + feature = "tos", + feature = "goosefs", +))] +pub(crate) mod opendal_store; pub mod providers; pub mod storage_options; #[cfg(test)] diff --git a/rust/lance-io/src/object_store/dynamic_opendal.rs b/rust/lance-io/src/object_store/dynamic_opendal.rs index ce9809be1d7..50b15d2180b 100644 --- a/rust/lance-io/src/object_store/dynamic_opendal.rs +++ b/rust/lance-io/src/object_store/dynamic_opendal.rs @@ -14,10 +14,10 @@ use object_store::{ ObjectStore as OSObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, }; -use object_store_opendal::OpendalStore; use tokio::sync::RwLock; use crate::object_store::StorageOptionsAccessor; +use crate::object_store::opendal_store::OpendalStore; use lance_core::Result; type NormalizeConfigFn = fn(&HashMap) -> Result>; diff --git a/rust/lance-io/src/object_store/opendal_store.rs b/rust/lance-io/src/object_store/opendal_store.rs new file mode 100644 index 00000000000..99523759095 --- /dev/null +++ b/rust/lance-io/src/object_store/opendal_store.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::fmt; +use std::ops::Range; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt, future, stream::BoxStream}; +use object_store::path::Path; +use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore as OSObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, + RenameOptions, +}; +use object_store_opendal::OpendalStore as InnerOpendalStore; +use opendal::Operator; + +/// Adapts OpenDAL listing paths to the spelling used by the request. +/// +/// The upstream bridge builds listed locations with [`Path::from`], which +/// percent-encodes reserved characters. Lance builds dataset base paths with +/// [`Path::from_url_path`], so mismatched listed locations must be decoded. +/// Locations that already match the requested prefix retain their spelling to +/// preserve paths containing literal percent escapes. +#[derive(Debug, Clone)] +pub(super) struct OpendalStore { + inner: InnerOpendalStore, +} + +impl OpendalStore { + pub(super) fn new(operator: Operator) -> Self { + Self { + inner: InnerOpendalStore::new(operator), + } + } +} + +impl fmt::Display for OpendalStore { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.inner.fmt(formatter) + } +} + +fn normalize_location(location: &Path, prefix: Option<&Path>) -> object_store::Result { + if prefix.is_none_or(|prefix| location.prefix_matches(prefix)) { + return Ok(location.clone()); + } + + Path::from_url_path(location.as_ref()).map_err(|source| object_store::Error::Generic { + store: "OpendalStore", + source: Box::new(source), + }) +} + +fn normalize_object_meta( + mut meta: ObjectMeta, + prefix: Option<&Path>, +) -> object_store::Result { + meta.location = normalize_location(&meta.location, prefix)?; + Ok(meta) +} + +#[async_trait] +impl OSObjectStore for OpendalStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + options: PutOptions, + ) -> object_store::Result { + self.inner.put_opts(location, payload, options).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + options: PutMultipartOptions, + ) -> object_store::Result> { + self.inner.put_multipart_opts(location, options).await + } + + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + self.inner.get_opts(location, options).await + } + + async fn get_ranges( + &self, + location: &Path, + ranges: &[Range], + ) -> object_store::Result> { + self.inner.get_ranges(location, ranges).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + self.inner.delete_stream(locations) + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result> { + let listed = self.inner.list(prefix); + let prefix = prefix.cloned(); + listed + .map(move |result| result.and_then(|meta| normalize_object_meta(meta, prefix.as_ref()))) + .boxed() + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, object_store::Result> { + if self.inner.info().capability().list_with_start_after { + let listed = self.inner.list_with_offset(prefix, offset); + let prefix = prefix.cloned(); + listed + .map(move |result| { + result.and_then(|meta| normalize_object_meta(meta, prefix.as_ref())) + }) + .boxed() + } else { + // The bridge's fallback compares its encoded output with the raw + // offset. Filter normalized locations so both sides use one form. + let offset = offset.clone(); + self.list(prefix) + .try_filter(move |meta| future::ready(meta.location > offset)) + .boxed() + } + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result { + let mut result = self.inner.list_with_delimiter(prefix).await?; + for object in &mut result.objects { + object.location = normalize_location(&object.location, prefix)?; + } + for common_prefix in &mut result.common_prefixes { + *common_prefix = normalize_location(common_prefix, prefix)?; + } + Ok(result) + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> object_store::Result<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + options: RenameOptions, + ) -> object_store::Result<()> { + self.inner.rename_opts(from, to, options).await + } +} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + use futures::TryStreamExt; + use object_store::ObjectStoreExt; + use opendal::services::Memory; + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::raw_reserved_character("tables/run~1/t.lance")] + #[case::literal_percent_escape("tables/run%25231/t.lance")] + #[tokio::test] + async fn test_list_preserves_request_path_spelling(#[case] base_url: &str) { + let operator = Operator::new(Memory::default()).unwrap(); + let store = OpendalStore::new(operator); + let base = Path::from_url_path(base_url).unwrap(); + let direct_location = base.clone().join("manifest.lance"); + let nested_location = Path::from_url_path(format!("{base_url}/data/part.lance")).unwrap(); + for location in [&direct_location, &nested_location] { + store + .put(location, Bytes::from_static(b"data").into()) + .await + .unwrap(); + } + + let listed = store + .list(Some(&base)) + .try_collect::>() + .await + .unwrap(); + let mut listed_locations = listed + .into_iter() + .map(|meta| meta.location) + .collect::>(); + listed_locations.sort(); + let mut expected_locations = vec![direct_location.clone(), nested_location.clone()]; + expected_locations.sort(); + assert_eq!(listed_locations, expected_locations); + assert!( + listed_locations + .iter() + .all(|location| location.prefix_matches(&base)) + ); + + let listed_after_nested = store + .list_with_offset(Some(&base), &nested_location) + .try_collect::>() + .await + .unwrap(); + assert_eq!(listed_after_nested.len(), 1); + assert_eq!(listed_after_nested[0].location, direct_location); + + let delimited = store.list_with_delimiter(Some(&base)).await.unwrap(); + assert_eq!(delimited.objects.len(), 1); + assert_eq!(delimited.objects[0].location, direct_location); + assert_eq!(delimited.common_prefixes, vec![base.clone().join("data")]); + } +} diff --git a/rust/lance-io/src/object_store/providers/aws.rs b/rust/lance-io/src/object_store/providers/aws.rs index 73708d1a05e..faea1d96a8d 100644 --- a/rust/lance-io/src/object_store/providers/aws.rs +++ b/rust/lance-io/src/object_store/providers/aws.rs @@ -10,7 +10,6 @@ use mock_instant::thread_local::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH}; use object_store::ObjectStore as OSObjectStore; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::S3}; use aws_config::default_provider::credentials::DefaultCredentialsChain; @@ -30,6 +29,7 @@ use object_store::{ use tokio::sync::RwLock; use url::Url; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, diff --git a/rust/lance-io/src/object_store/providers/azure.rs b/rust/lance-io/src/object_store/providers/azure.rs index eb547fa9c3d..5298a592138 100644 --- a/rust/lance-io/src/object_store/providers/azure.rs +++ b/rust/lance-io/src/object_store/providers/azure.rs @@ -9,7 +9,6 @@ use std::{ }; use object_store::ObjectStore as OSObjectStore; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::Azblob, services::Azdls}; use object_store::{ @@ -18,6 +17,7 @@ use object_store::{ }; use url::Url; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, diff --git a/rust/lance-io/src/object_store/providers/gcp.rs b/rust/lance-io/src/object_store/providers/gcp.rs index 2a1570b6a13..8dd45446112 100644 --- a/rust/lance-io/src/object_store/providers/gcp.rs +++ b/rust/lance-io/src/object_store/providers/gcp.rs @@ -7,7 +7,6 @@ use object_store::{ ClientOptions, CredentialProvider, ObjectStore as OSObjectStore, Result as ObjectStoreResult, client::{HttpClient, HttpConnector, HttpRequestBody, ReqwestConnector}, }; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::Gcs}; use reqsign_core::{Context as ReqsignContext, HttpSend, OsEnv, ProvideCredential}; use reqsign_file_read_tokio::TokioFileRead; @@ -20,6 +19,7 @@ use object_store::{ }; use url::Url; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor, diff --git a/rust/lance-io/src/object_store/providers/goosefs.rs b/rust/lance-io/src/object_store/providers/goosefs.rs index 72b64343169..375b8a6615e 100644 --- a/rust/lance-io/src/object_store/providers/goosefs.rs +++ b/rust/lance-io/src/object_store/providers/goosefs.rs @@ -4,10 +4,10 @@ use std::collections::HashMap; use std::sync::Arc; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::GooseFs}; use url::Url; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, diff --git a/rust/lance-io/src/object_store/providers/huggingface.rs b/rust/lance-io/src/object_store/providers/huggingface.rs index 0b0b6bc8f97..205e22273ca 100644 --- a/rust/lance-io/src/object_store/providers/huggingface.rs +++ b/rust/lance-io/src/object_store/providers/huggingface.rs @@ -6,11 +6,11 @@ use std::sync::Arc; use object_store::ObjectStore as OSObjectStore; use object_store::path::Path; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::Huggingface}; use url::Url; use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::parse_hf_repo_id; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, diff --git a/rust/lance-io/src/object_store/providers/oss.rs b/rust/lance-io/src/object_store/providers/oss.rs index 4a97547721f..0e45cf5c5bf 100644 --- a/rust/lance-io/src/object_store/providers/oss.rs +++ b/rust/lance-io/src/object_store/providers/oss.rs @@ -5,11 +5,11 @@ use std::collections::HashMap; use std::sync::Arc; use object_store::ObjectStore as OSObjectStore; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::Oss}; use url::Url; use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, diff --git a/rust/lance-io/src/object_store/providers/tencent.rs b/rust/lance-io/src/object_store/providers/tencent.rs index 8557ee70f21..481eadf621a 100644 --- a/rust/lance-io/src/object_store/providers/tencent.rs +++ b/rust/lance-io/src/object_store/providers/tencent.rs @@ -4,10 +4,10 @@ use std::collections::HashMap; use std::sync::Arc; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::Cos}; use url::Url; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, diff --git a/rust/lance-io/src/object_store/providers/tos.rs b/rust/lance-io/src/object_store/providers/tos.rs index b53b30b997f..3c9fa88c4eb 100644 --- a/rust/lance-io/src/object_store/providers/tos.rs +++ b/rust/lance-io/src/object_store/providers/tos.rs @@ -5,11 +5,11 @@ use std::collections::HashMap; use std::sync::Arc; use object_store::ObjectStore as OSObjectStore; -use object_store_opendal::OpendalStore; use opendal::{Operator, services::Tos}; use url::Url; use crate::object_store::dynamic_opendal::DynamicOpenDalStore; +use crate::object_store::opendal_store::OpendalStore; use crate::object_store::{ DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions,