1use std::path::{Path, PathBuf};
4
5#[derive(Clone, Debug, Eq, PartialEq)]
6pub enum InputSource {
7 Local(PathBuf),
8 S3(String),
9 Gcs(String),
10 Azure(String),
13 Http(String),
14}
15
16pub fn input_source(path: &Path) -> InputSource {
18 let s = path.as_os_str().to_string_lossy();
19 if let Some(after_scheme) = s.find("://") {
20 let prefix = s[..after_scheme].to_lowercase();
21 let rest = s[after_scheme + 3..].to_string();
22 if prefix == "s3" || prefix == "s3a" {
23 return InputSource::S3(rest);
24 }
25 if prefix == "gs" || prefix == "gcs" {
26 return InputSource::Gcs(rest);
27 }
28 if let Some((account, container, key)) = azure_parts(&s) {
29 return InputSource::Azure(azure_url(&account, &container, &key));
30 }
31 if prefix == "http" || prefix == "https" {
32 return InputSource::Http(s.to_string());
33 }
34 }
35 InputSource::Local(path.to_path_buf())
36}
37
38pub fn is_remote_url(path: &Path) -> bool {
42 !matches!(input_source(path), InputSource::Local(_))
43 || path
44 .to_string_lossy()
45 .split_once("://")
46 .is_some_and(|(scheme, _)| is_azure_short_scheme(scheme))
47}
48
49pub fn opens_in_this_build(path: &Path) -> bool {
52 match input_source(path) {
53 InputSource::Local(_) if is_remote_url(path) => cfg!(feature = "cloud"),
54 InputSource::Local(_) => true,
55 InputSource::S3(_) | InputSource::Gcs(_) | InputSource::Azure(_) => {
56 cfg!(feature = "cloud")
57 }
58 InputSource::Http(_) => cfg!(feature = "http"),
59 }
60}
61
62pub(crate) fn is_azure_short_scheme(scheme: &str) -> bool {
64 matches!(scheme.to_ascii_lowercase().as_str(), "az" | "adl" | "azure")
65}
66
67pub fn split_source_id(url: &str) -> (Option<&str>, std::borrow::Cow<'_, str>) {
73 let Some((scheme, rest)) = url.split_once("://") else {
74 return (None, url.into());
75 };
76 if !matches!(scheme.to_ascii_lowercase().as_str(), "s3" | "s3a") {
77 return (None, url.into());
78 }
79 let first = rest.split('/').next().unwrap_or(rest);
80 match first.split_once('@') {
81 Some((id, _)) if !id.is_empty() => {
82 let plain = format!("{scheme}://{}", &rest[id.len() + 1..]);
83 (Some(id), plain.into())
84 }
85 _ => (None, url.into()),
86 }
87}
88
89pub fn azure_parts(url: &str) -> Option<(String, String, String)> {
97 let (scheme, rest) = url.split_once("://")?;
98 let scheme = scheme.to_ascii_lowercase();
99 let (host_part, path) = match rest.split_once('/') {
100 Some((host, path)) => (host, path),
101 None => (rest, ""),
102 };
103 let account_of = |host: &str| {
104 let host = host.to_ascii_lowercase();
105 [".dfs.core.windows.net", ".blob.core.windows.net"]
106 .iter()
107 .find_map(|suffix| host.strip_suffix(suffix).map(str::to_string))
108 .filter(|account| !account.is_empty() && !account.contains('.'))
109 };
110 match scheme.as_str() {
111 "abfss" | "abfs" => {
112 let (container, host) = host_part.split_once('@')?;
113 let account = account_of(host)?;
114 (!container.is_empty()).then(|| (account, container.to_string(), path.to_string()))
115 }
116 "https" | "http" => {
117 let account = account_of(host_part)?;
118 let (container, path) = match path.split_once('/') {
119 Some((container, path)) => (container, path),
120 None => (path, ""),
121 };
122 (!container.is_empty()).then(|| (account, container.to_string(), path.to_string()))
123 }
124 _ => None,
125 }
126}
127
128pub fn azure_url(account: &str, container: &str, path: &str) -> String {
130 format!(
131 "abfss://{container}@{account}.dfs.core.windows.net/{}",
132 path.trim_start_matches('/')
133 )
134}
135
136pub(crate) fn canonical_cloud_place(url: &str) -> String {
139 let canonical = match azure_parts(url) {
140 Some((account, container, path)) => azure_url(&account, &container, &path),
141 None => url.to_string(),
142 };
143 canonical.trim_end_matches('/').to_string()
144}
145
146pub(crate) fn is_prefix_or_glob(url: &str) -> bool {
148 url.ends_with('/') || url.contains('*')
149}
150
151pub(crate) fn has_glob_chars(path: &Path) -> bool {
153 path.as_os_str().to_string_lossy().contains(['*', '?', '['])
154}
155
156pub(crate) fn expands_as_glob(path: &Path) -> bool {
161 has_glob_chars(path) && std::fs::symlink_metadata(path).is_err()
162}
163
164pub(crate) fn polars_literal_path(
168 path: &Path,
169) -> polars::prelude::PolarsResult<polars::prelude::PlRefPath> {
170 if !has_glob_chars(path) || expands_as_glob(path) {
171 return polars::prelude::PlRefPath::try_from_path(path);
172 }
173 let text = polars::prelude::PlRefPath::try_from_path(path)?;
176 Ok(polars::prelude::PlRefPath::new(
177 escape_glob(text.as_str()).as_str(),
178 ))
179}
180
181pub(crate) fn escape_glob(text: &str) -> String {
183 let mut escaped = String::with_capacity(text.len() + 8);
184 for c in text.chars() {
185 if matches!(c, '*' | '?' | '[' | ']') {
186 escaped.extend(['[', c, ']']);
187 } else {
188 escaped.push(c);
189 }
190 }
191 escaped
192}
193
194pub(crate) fn scans_in_place(path: &Path) -> bool {
199 if !matches!(
200 input_source(path),
201 InputSource::S3(_) | InputSource::Gcs(_) | InputSource::Azure(_)
202 ) {
203 return false;
204 }
205 let url = path.to_string_lossy();
206 let (_, ext) = url_path_extension(&url);
207 !cloud_path_should_download(ext.as_deref(), is_prefix_or_glob(&url))
208}
209
210pub(crate) fn url_path_extension(url: &str) -> (String, Option<String>) {
213 let path_part = if let Some(i) = url.find("://") {
214 let scheme = url[..i].to_lowercase();
215 let after = &url[i + 3..];
216 if scheme == "http" || scheme == "https" {
217 after
218 .find('/')
219 .map(|j| after[j + 1..].to_string())
220 .unwrap_or_default()
221 } else {
222 after.to_string()
223 }
224 } else {
225 String::new()
226 };
227 let last_segment = path_part.rsplit('/').next().unwrap_or(&path_part);
228 let ext = std::path::Path::new(last_segment)
229 .extension()
230 .and_then(|e| e.to_str())
231 .map(String::from);
232 (path_part, ext)
233}
234
235#[cfg(any(feature = "http", feature = "cloud"))]
239pub(crate) fn download_suffix(url: &str) -> Option<String> {
240 let (path_part, ext) = url_path_extension(url);
241 let ext = ext?;
242 const COMPRESSION: [&str; 6] = ["gz", "zst", "bz2", "xz", "lz4", "zip"];
243 if !COMPRESSION.iter().any(|c| ext.eq_ignore_ascii_case(c)) {
244 return Some(ext);
245 }
246 let name = path_part.rsplit('/').next().unwrap_or(&path_part);
247 let stem = &name[..name.len() - ext.len() - 1];
248 match Path::new(stem).extension().and_then(|e| e.to_str()) {
249 Some(inner) => Some(format!("{inner}.{ext}")),
250 None => Some(ext),
251 }
252}
253
254pub(crate) fn cloud_path_should_download(ext: Option<&str>, is_glob: bool) -> bool {
259 if is_glob {
260 return false;
261 }
262 match ext {
263 None => false,
264 Some(e) => !e.eq_ignore_ascii_case("parquet"),
265 }
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271
272 #[test]
273 fn an_existing_name_is_never_a_glob() {
274 let dir = tempfile::tempdir().unwrap();
275 let file = dir.path().join("d[1].csv");
276 std::fs::write(&file, "a\n1\n").unwrap();
277 assert!(!expands_as_glob(&file));
278 assert!(expands_as_glob(&dir.path().join("d[2].csv")));
279 assert!(expands_as_glob(&dir.path().join("*.csv")));
280 assert!(!expands_as_glob(&dir.path().join("plain.csv")));
281 assert!(!expands_as_glob(dir.path()));
282
283 let escaped = polars_literal_path(&file).unwrap();
284 assert!(escaped.as_str().ends_with("d[[]1[]].csv"), "{escaped:?}");
285 let pattern = dir.path().join("d[2].csv");
286 let kept = polars_literal_path(&pattern).unwrap();
287 assert_eq!(
289 kept.as_str().replace('\\', "/"),
290 pattern.to_str().unwrap().replace('\\', "/")
291 );
292 }
293
294 #[test]
298 fn an_escaped_name_reads_that_file_alone() {
299 use polars::prelude::{LazyFileListReader, LazyJsonLineReader};
300 let dir = tempfile::tempdir().unwrap();
301 let nested = dir.path().join("set[1]");
302 std::fs::create_dir(&nested).unwrap();
303 std::fs::create_dir(dir.path().join("set1")).unwrap();
304 let mut names = vec!["d[1]", "h]", "i{j,k}", "[!x]", "set[1]/p[a]"];
305 if cfg!(unix) {
306 names.extend(["a*b", "x?", "e\\f", "g[x]*?"]);
307 }
308 for decoy in ["d1", "ha", "set1/pa", "ab", "xy", "gx", "y"] {
310 std::fs::write(dir.path().join(format!("{decoy}.jsonl")), "{\"v\": 0}\n").unwrap();
311 }
312 for name in names {
313 let file = dir.path().join(format!("{name}.jsonl"));
314 std::fs::write(
315 &file,
316 format!("{{\"v\": \"{}\"}}\n", name.replace('\\', "/")),
317 )
318 .unwrap();
319 let lf = LazyJsonLineReader::new(polars_literal_path(&file).unwrap())
320 .finish()
321 .unwrap();
322 let df = lf.collect().unwrap_or_else(|e| panic!("{name}: {e}"));
323 assert_eq!(df.height(), 1, "{name}");
324 let v = df
325 .column("v")
326 .unwrap()
327 .str()
328 .unwrap()
329 .get(0)
330 .map(str::to_string);
331 assert_eq!(
332 v.as_deref(),
333 Some(name.replace('\\', "/").as_str()),
334 "{name}"
335 );
336 }
337 }
338
339 #[test]
340 fn every_url_datui_reads_is_remote() {
341 for url in [
342 "s3://bucket/key.parquet",
343 "s3a://bucket/key.parquet",
344 "gs://bucket/dir/",
345 "gcs://bucket/dir/",
346 "abfss://release@overturemapswestus2.dfs.core.windows.net/2026-09-23.1/",
347 "abfs://container@account.dfs.core.windows.net/x.parquet",
348 "https://account.blob.core.windows.net/container/x.csv",
349 "az://container/x.csv",
350 "adl://container/x.csv",
351 "azure://container/x.csv",
352 "http://example.com/data.csv",
353 "https://example.com/data.csv",
354 ] {
355 assert!(is_remote_url(Path::new(url)), "{url}");
356 }
357 for path in ["/tmp/file.parquet", "relative.csv", ".", "data/2024.csv"] {
358 assert!(!is_remote_url(Path::new(path)), "{path}");
359 }
360 }
361
362 #[test]
363 fn input_source_local_path() {
364 let p = PathBuf::from("/tmp/file.parquet");
365 assert!(matches!(input_source(&p), InputSource::Local(_)));
366 let p = PathBuf::from("relative.csv");
367 assert!(matches!(input_source(&p), InputSource::Local(_)));
368 let p = PathBuf::from(".");
369 assert!(matches!(input_source(&p), InputSource::Local(_)));
370 }
371
372 #[test]
373 fn input_source_s3() {
374 let p = PathBuf::from("s3://bucket/key.parquet");
375 match input_source(&p) {
376 InputSource::S3(rest) => assert_eq!(rest, "bucket/key.parquet"),
377 _ => panic!("expected S3"),
378 }
379 let p = PathBuf::from("S3://my-bucket/path/to/file.csv");
380 match input_source(&p) {
381 InputSource::S3(rest) => assert_eq!(rest, "my-bucket/path/to/file.csv"),
382 _ => panic!("expected S3"),
383 }
384 }
385
386 #[test]
387 fn a_source_is_split_off_s3_urls_only() {
388 assert_eq!(
389 split_source_id("s3://onprem@sales/2024/q3.parquet"),
390 (Some("onprem"), "s3://sales/2024/q3.parquet".into())
391 );
392 assert_eq!(
393 split_source_id("s3://onprem@sales"),
394 (Some("onprem"), "s3://sales".into())
395 );
396 assert_eq!(
397 split_source_id("s3://sales/a@b.parquet"),
398 (None, "s3://sales/a@b.parquet".into())
399 );
400 assert_eq!(
401 split_source_id("gs://bucket/key"),
402 (None, "gs://bucket/key".into())
403 );
404 assert_eq!(
405 split_source_id("https://user@host/file.csv"),
406 (None, "https://user@host/file.csv".into())
407 );
408 assert_eq!(split_source_id("s3://@sales"), (None, "s3://@sales".into()));
409 }
410
411 #[test]
412 fn azure_urls_in_every_form_that_names_the_account_become_one() {
413 let canonical = "abfss://datui-test@datalake001.dfs.core.windows.net/demo/fred/";
414 for url in [
415 canonical,
416 "abfs://datui-test@datalake001.dfs.core.windows.net/demo/fred/",
417 "https://datalake001.blob.core.windows.net/datui-test/demo/fred/",
418 "https://DataLake001.dfs.core.windows.net/datui-test/demo/fred/",
419 ] {
420 assert_eq!(
421 input_source(Path::new(url)),
422 InputSource::Azure(canonical.to_string()),
423 "{url}"
424 );
425 }
426 assert_eq!(
427 azure_parts("abfss://c@acct.dfs.core.windows.net"),
428 Some(("acct".to_string(), "c".to_string(), String::new()))
429 );
430 assert_eq!(azure_parts("az://container/path"), None);
432 assert!(matches!(
433 input_source(Path::new("https://example.com/c/x.csv")),
434 InputSource::Http(_)
435 ));
436 assert!(scans_in_place(Path::new(
437 "abfss://c@acct.dfs.core.windows.net/x.parquet"
438 )));
439 }
440
441 #[test]
442 fn input_source_http() {
443 let p = PathBuf::from("https://example.com/data.parquet");
444 match input_source(&p) {
445 InputSource::Http(u) => assert_eq!(u, "https://example.com/data.parquet"),
446 _ => panic!("expected Http"),
447 }
448 let p = PathBuf::from("http://host/path/file.csv");
449 match input_source(&p) {
450 InputSource::Http(u) => assert_eq!(u, "http://host/path/file.csv"),
451 _ => panic!("expected Http"),
452 }
453 }
454
455 #[test]
456 fn input_source_gcs() {
457 let p = PathBuf::from("gs://my-bucket/path/file.parquet");
458 match input_source(&p) {
459 InputSource::Gcs(rest) => assert_eq!(rest, "my-bucket/path/file.parquet"),
460 _ => panic!("expected Gcs"),
461 }
462 let p = PathBuf::from("gcs://bucket/key.parquet");
463 match input_source(&p) {
464 InputSource::Gcs(rest) => assert_eq!(rest, "bucket/key.parquet"),
465 _ => panic!("expected Gcs"),
466 }
467 }
468
469 #[test]
470 fn input_source_unknown_scheme_stays_local() {
471 let p = PathBuf::from("file:///tmp/foo.parquet");
472 assert!(matches!(input_source(&p), InputSource::Local(_)));
473 }
474
475 #[test]
476 fn url_path_extension_s3() {
477 let (path, ext) = url_path_extension("s3://bucket/key.parquet");
478 assert_eq!(path, "bucket/key.parquet");
479 assert_eq!(ext.as_deref(), Some("parquet"));
480 let (path, ext) = url_path_extension("s3://b/path/to/file.csv");
481 assert_eq!(path, "b/path/to/file.csv");
482 assert_eq!(ext.as_deref(), Some("csv"));
483 }
484
485 #[cfg(any(feature = "http", feature = "cloud"))]
486 #[test]
487 fn a_download_keeps_what_the_compressed_file_holds() {
488 assert_eq!(
489 download_suffix("s3://b/edge/100%.csv.gz").as_deref(),
490 Some("csv.gz")
491 );
492 assert_eq!(
493 download_suffix("https://x.com/a/log.json.zst").as_deref(),
494 Some("json.zst")
495 );
496 assert_eq!(download_suffix("gs://b/data.csv").as_deref(), Some("csv"));
497 assert_eq!(download_suffix("s3://b/archive.gz").as_deref(), Some("gz"));
498 assert_eq!(download_suffix("s3://b/no-extension"), None);
499 }
500
501 #[test]
502 fn url_path_extension_https() {
503 let (path, ext) = url_path_extension("https://example.com/dir/file.parquet");
504 assert_eq!(path, "dir/file.parquet");
505 assert_eq!(ext.as_deref(), Some("parquet"));
506 let (_, ext) = url_path_extension("https://x.com/file.csv.gz");
507 assert_eq!(ext.as_deref(), Some("gz"));
508 }
509
510 #[test]
511 fn only_parquet_prefixes_and_globs_are_scanned_in_place() {
512 assert!(scans_in_place(Path::new("s3://bucket/obj.parquet")));
513 assert!(scans_in_place(Path::new("gs://bucket/prefix/")));
514 assert!(scans_in_place(Path::new("s3://bucket/year=*/*.parquet")));
515 assert!(!scans_in_place(Path::new("s3://bucket/data.csv")));
517 assert!(!scans_in_place(Path::new("s3://bucket/data.csv.gz")));
518 assert!(!scans_in_place(Path::new(
519 "https://example.com/data.parquet"
520 )));
521 assert!(!scans_in_place(Path::new("/data/local.parquet")));
522 }
523
524 #[test]
525 fn cloud_path_should_download() {
526 assert!(super::cloud_path_should_download(Some("csv"), false));
527 assert!(super::cloud_path_should_download(Some("gz"), false));
528 assert!(super::cloud_path_should_download(Some("csv.gz"), false));
529 assert!(!super::cloud_path_should_download(Some("parquet"), false));
530 assert!(!super::cloud_path_should_download(None, false));
531 assert!(!super::cloud_path_should_download(Some("csv"), true));
532 assert!(!super::cloud_path_should_download(Some("parquet"), true));
533 }
534}