feat: Support S3 virtual-hosted–style URI (#24405)

Author: nameexhaustion

crates/polars-io/src/path_utils/hugging_face.rs

@@ -215,7 +215,7 @@ impl GetPages<'_> {
 pub(super) async fn expand_paths_hf(
     paths: &[PlPath],
     check_directory_level: bool,
-    cloud_options: Option<&CloudOptions>,
+    cloud_options: &Option<CloudOptions>,
     glob: bool,
 ) -> PolarsResult<(usize, Vec<PlPath>)> {
     assert!(!paths.is_empty());

crates/polars-io/src/path_utils/mod.rs

@@ -163,7 +163,7 @@ pub fn expanded_from_single_directory(addrs: &[PlPath], expanded_addrs: &[PlPath
 pub fn expand_paths(
     paths: &[PlPath],
     glob: bool,
-    #[allow(unused_variables)] cloud_options: Option<&CloudOptions>,
+    #[allow(unused_variables)] cloud_options: &mut Option<CloudOptions>,
 ) -> PolarsResult<Arc<[PlPath]>> {
     expand_paths_hive(paths, glob, cloud_options, false).map(|x| x.0)
 }
@@ -204,7 +204,7 @@ impl HiveIdxTracker<'_> {
 pub fn expand_paths_hive(
     paths: &[PlPath],
     glob: bool,
-    #[allow(unused_variables)] cloud_options: Option<&CloudOptions>,
+    #[allow(unused_variables)] cloud_options: &mut Option<CloudOptions>,
     check_directory_level: bool,
 ) -> PolarsResult<(Arc<[PlPath]>, usize)> {
     let Some(first_path) = paths.first() else {
@@ -398,13 +398,65 @@ pub fn expand_paths_hive(
             };
 
             for (path_idx, path) in paths.iter().enumerate() {
+                use std::borrow::Cow;
+
+                let mut path = Cow::Borrowed(path);
+
                 if matches!(
                     path.cloud_scheme(),
                     Some(CloudScheme::Http | CloudScheme::Https)
                 ) {
-                    out_paths.push(path.clone());
-                    hive_idx_tracker.update(0, path_idx)?;
-                    continue;
+                    let mut rewrite_aws = false;
+
+                    #[cfg(feature = "aws")]
+                    {
+                        // See https://docs.aws.amazon.com/AmazonS3/latest/userguide/VirtualHosting.html#virtual-hosted-style-access
+                        // Path format: https://bucket-name.s3.region-code.amazonaws.com/key-name
+                        let p = path.as_ref().as_ref();
+                        let after_scheme = p.strip_scheme();
+                        if let Some(bucket_end) = after_scheme.find(".s3.") {
+                            if let Some(region_end) = after_scheme.find(".amazonaws.com/") {
+                                if bucket_end < region_end
+                                    && region_end < after_scheme.find("/").unwrap()
+                                {
+                                    use crate::cloud::CloudConfig;
+
+                                    rewrite_aws = true;
+
+                                    let bucket = &after_scheme[..bucket_end];
+                                    let region = &after_scheme[bucket_end + 4..region_end];
+                                    let key = &after_scheme[region_end + 15..];
+
+                                    if let CloudConfig::Aws(configs) = cloud_options
+                                        .get_or_insert_default()
+                                        .config
+                                        .get_or_insert_with(|| {
+                                            CloudConfig::Aws(Vec::with_capacity(1))
+                                        })
+                                    {
+                                        use object_store::aws::AmazonS3ConfigKey;
+
+                                        if !matches!(
+                                            configs.last(),
+                                            Some((AmazonS3ConfigKey::Region, _))
+                                        ) {
+                                            configs.push((AmazonS3ConfigKey::Region, region.into()))
+                                        }
+                                    }
+
+                                    path = Cow::Owned(PlPath::from_string(format!(
+                                        "s3://{bucket}/{key}"
+                                    )))
+                                }
+                            }
+                        }
+                    }
+
+                    if !rewrite_aws {
+                        out_paths.push(path.into_owned());
+                        hive_idx_tracker.update(0, path_idx)?;
+                        continue;
+                    }
                 }
 
                 let glob_start_idx = get_glob_start_idx(path.to_str().as_bytes());
@@ -413,7 +465,7 @@ pub fn expand_paths_hive(
                     path.clone()
                 } else {
                     let (expand_start_idx, paths) =
-                        expand_path_cloud(path.to_str(), cloud_options)?;
+                        expand_path_cloud(path.to_str(), cloud_options.as_ref())?;
                     out_paths.extend_from_slice(&paths);
                     hive_idx_tracker.update(expand_start_idx, path_idx)?;
                     continue;
@@ -422,7 +474,7 @@ pub fn expand_paths_hive(
                 hive_idx_tracker.update(0, path_idx)?;
 
                 let iter = crate::pl_async::get_runtime()
-                    .block_in_place_on(crate::async_glob(path.to_str(), cloud_options))?;
+                    .block_in_place_on(crate::async_glob(path.to_str(), cloud_options.as_ref()))?;
 
                 if is_cloud {
                     out_paths.extend(iter.into_iter().map(PlPath::from_string));
@@ -577,7 +629,7 @@ mod tests {
 
         let path = "https://pola.rs/test.csv?token=bear";
         let paths = &[PlPath::new(path)];
-        let out = expand_paths(paths, true, None).unwrap();
+        let out = expand_paths(paths, true, &mut None).unwrap();
         assert_eq!(out.as_ref(), paths);
     }
 }

crates/polars-lazy/src/scan/csv.rs

@@ -275,7 +275,7 @@ impl LazyCsvReader {
             ScanSources::Paths(paths) => {
                 // TODO: Path expansion should happen when converting to the IR
                 // https://github.com/pola-rs/polars/issues/17634
-                let paths = expand_paths(&paths[..], self.glob(), self.cloud_options())?;
+                let paths = expand_paths(&paths[..], self.glob(), &mut self.cloud_options)?;
 
                 let Some(path) = paths.first() else {
                     polars_bail!(ComputeError: "no paths specified for this reader");

crates/polars-plan/src/dsl/scan_sources.rs

@@ -159,16 +159,12 @@ impl PartialEq for ScanSources {
 impl Eq for ScanSources {}
 
 impl ScanSources {
-    pub fn expand_paths(
-        &self,
-        scan_args: &UnifiedScanArgs,
-        #[allow(unused_variables)] cloud_options: Option<&CloudOptions>,
-    ) -> PolarsResult<Self> {
+    pub fn expand_paths(&self, scan_args: &mut UnifiedScanArgs) -> PolarsResult<Self> {
         match self {
             Self::Paths(paths) => Ok(Self::Paths(expand_paths(
                 paths,
                 scan_args.glob,
-                cloud_options,
+                &mut scan_args.cloud_options,
             )?)),
             v => Ok(v.clone()),
         }
@@ -180,14 +176,13 @@ impl ScanSources {
     pub fn expand_paths_with_hive_update(
         &self,
         scan_args: &mut UnifiedScanArgs,
-        #[allow(unused_variables)] cloud_options: Option<&CloudOptions>,
     ) -> PolarsResult<Self> {
         match self {
             Self::Paths(paths) => {
                 let (expanded_paths, hive_start_idx) = expand_paths_hive(
                     paths,
                     scan_args.glob,
-                    cloud_options,
+                    &mut scan_args.cloud_options,
                     scan_args.hive_options.enabled.unwrap_or(false),
                 )?;
 

crates/polars-plan/src/plans/conversion/dsl_to_ir/scans.rs

@@ -1,6 +1,5 @@
 use either::Either;
 use polars_io::RowIndex;
-use polars_io::cloud::CloudOptions;
 #[cfg(feature = "cloud")]
 use polars_io::pl_async::get_runtime;
 use polars_io::prelude::*;
@@ -20,9 +19,6 @@ pub(super) fn dsl_to_ir(
     let mut cached_ir = cached_ir.lock().unwrap();
 
     if cached_ir.is_none() {
-        let cloud_options = unified_scan_args_box.cloud_options.clone();
-        let cloud_options = cloud_options.as_ref();
-
         let unified_scan_args = unified_scan_args_box.as_mut();
 
         if let Some(hive_schema) = unified_scan_args.hive_options.schema.as_deref() {
@@ -45,16 +41,14 @@ pub(super) fn dsl_to_ir(
         let sources = match &*scan_type {
             #[cfg(feature = "parquet")]
             FileScanDsl::Parquet { .. } => {
-                sources.expand_paths_with_hive_update(unified_scan_args, cloud_options)?
+                sources.expand_paths_with_hive_update(unified_scan_args)?
             },
             #[cfg(feature = "ipc")]
-            FileScanDsl::Ipc { .. } => {
-                sources.expand_paths_with_hive_update(unified_scan_args, cloud_options)?
-            },
+            FileScanDsl::Ipc { .. } => sources.expand_paths_with_hive_update(unified_scan_args)?,
             #[cfg(feature = "csv")]
-            FileScanDsl::Csv { .. } => sources.expand_paths(unified_scan_args, cloud_options)?,
+            FileScanDsl::Csv { .. } => sources.expand_paths(unified_scan_args)?,
             #[cfg(feature = "json")]
-            FileScanDsl::NDJson { .. } => sources.expand_paths(unified_scan_args, cloud_options)?,
+            FileScanDsl::NDJson { .. } => sources.expand_paths(unified_scan_args)?,
             #[cfg(feature = "python")]
             FileScanDsl::PythonDataset { .. } => {
                 // There are a lot of places that short-circuit if the paths is empty,
@@ -71,7 +65,6 @@ pub(super) fn dsl_to_ir(
             &sources,
             sources_before_expansion,
             unified_scan_args,
-            cloud_options,
             ctxt.verbose,
         )?;
 
@@ -522,7 +515,6 @@ impl SourcesToFileInfo {
         sources: &ScanSources,
         sources_before_expansion: &ScanSources,
         unified_scan_args: &mut UnifiedScanArgs,
-        cloud_options: Option<&CloudOptions>,
     ) -> PolarsResult<(FileInfo, FileScanIR)> {
         let require_first_source = |failed_operation_name: &'static str, hint: &'static str| {
             sources.first_or_empty_expand_err(
@@ -533,6 +525,8 @@ impl SourcesToFileInfo {
             )
         };
 
+        let cloud_options = unified_scan_args.cloud_options.as_ref();
+
         Ok(match scan_type {
             #[cfg(feature = "parquet")]
             FileScanDsl::Parquet { options } => {
@@ -666,7 +660,6 @@ this scan to succeed with an empty DataFrame.",
         sources: &ScanSources,
         sources_before_expansion: &ScanSources,
         unified_scan_args: &mut UnifiedScanArgs,
-        cloud_options: Option<&CloudOptions>,
         verbose: bool,
     ) -> PolarsResult<(FileInfo, FileScanIR)> {
         // Only cache non-empty paths. Others are directly parsed.
@@ -679,7 +672,6 @@ this scan to succeed with an empty DataFrame.",
                     sources,
                     sources_before_expansion,
                     unified_scan_args,
-                    cloud_options,
                 );
             },
         };
@@ -731,7 +723,6 @@ this scan to succeed with an empty DataFrame.",
                     sources,
                     sources_before_expansion,
                     unified_scan_args,
-                    cloud_options,
                 );
             },
         };
@@ -747,7 +738,6 @@ this scan to succeed with an empty DataFrame.",
                 sources,
                 sources_before_expansion,
                 unified_scan_args,
-                cloud_options,
             )?;
             self.inner.insert(k, v.clone());
             Ok(v)

py-polars/polars/io/cloud/_utils.py

@@ -53,8 +53,14 @@ def _get_path_scheme(path: str | Path) -> str | None:
     return path_str[:i] if i >= 0 else None
 
 
-def _is_aws_cloud(scheme: str) -> bool:
-    return any(scheme == x for x in ["s3", "s3a"])
+def _is_aws_cloud(*, scheme: str, first_scan_path: str) -> bool:
+    return any(scheme == x for x in ["s3", "s3a"]) or (
+        (scheme == "http" or scheme == "https")
+        and 0
+        < first_scan_path.find(".s3.")
+        < first_scan_path.find(".amazonaws.com/")
+        < first_scan_path[len(scheme) + 3 :].find("/")
+    )
 
 
 def _is_azure_cloud(scheme: str) -> bool:

py-polars/polars/io/cloud/credential_provider/_builder.py

@@ -337,10 +337,10 @@ def f() -> CredentialProviderBuilder | None:
                 credential_provider
             )
 
-        if (path := _first_scan_path(source)) is None:
+        if (first_scan_path := _first_scan_path(source)) is None:
             return None
 
-        if (scheme := _get_path_scheme(path)) is None:
+        if (scheme := _get_path_scheme(first_scan_path)) is None:
             return None
 
         if _is_azure_cloud(scheme):
@@ -374,7 +374,9 @@ def f() -> CredentialProviderBuilder | None:
 
             storage_account = (
                 # Prefer the one embedded in the path
-                CredentialProviderAzure._extract_adls_uri_storage_account(str(path))
+                CredentialProviderAzure._extract_adls_uri_storage_account(
+                    str(first_scan_path)
+                )
                 or storage_account
             )
 
@@ -386,7 +388,7 @@ def f() -> CredentialProviderBuilder | None:
                 )
             )
 
-        elif _is_aws_cloud(scheme):
+        elif _is_aws_cloud(scheme=scheme, first_scan_path=str(first_scan_path)):
             region = None
             profile = None
             default_region = None