query_planning/source/
location.rs

1//! Source location and read-capability contracts.
2
3use std::{
4    fmt,
5    path::{Path, PathBuf},
6};
7
8use super::diagnostics::{object_store_diagnostic_id, redacted_debug_value, redacted_stable_id_diagnostic};
9use crate::util::fnv1a64;
10
11/// Explicit location/read-capability model for native Parquet sources.
12#[derive(Clone, PartialEq, Eq, Hash)]
13pub enum ParquetSourceLocation {
14    /// Local filesystem path readable by cuDF local Parquet readers.
15    LocalPath(PathBuf),
16    /// Object-store URI plus parsed scheme.
17    ObjectStoreUri {
18        /// Lowercase object-store scheme such as `s3`.
19        scheme: String,
20        /// Original object-store URI.
21        uri: String,
22    },
23}
24
25impl fmt::Debug for ParquetSourceLocation {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        match self {
28            Self::LocalPath(path) => f
29                .debug_tuple("LocalPath")
30                .field(&redacted_debug_value("path", path.display().to_string().as_str()))
31                .finish(),
32            Self::ObjectStoreUri { scheme, uri } => f
33                .debug_struct("ObjectStoreUri")
34                .field("scheme", scheme)
35                .field("uri", &redacted_debug_value("uri", uri))
36                .finish(),
37        }
38    }
39}
40
41impl ParquetSourceLocation {
42    /// Build a local filesystem source location.
43    pub fn local_path(path: impl Into<PathBuf>) -> Self {
44        Self::LocalPath(path.into())
45    }
46
47    /// Build a supported object-store URI location.
48    ///
49    /// Currently only `s3://...` is accepted. Errors are stable reason strings
50    /// such as `missing_uri_scheme` or `unsupported_uri_scheme:<scheme>`.
51    pub fn object_store_uri(uri: impl Into<String>) -> Result<Self, String> {
52        let uri = uri.into();
53        let Some((scheme, _)) = uri.split_once("://") else {
54            return Err("missing_uri_scheme".to_string());
55        };
56        if scheme.is_empty() {
57            return Err("empty_uri_scheme".to_string());
58        }
59        let scheme = scheme.to_ascii_lowercase();
60        if scheme != "s3" {
61            return Err(format!("unsupported_uri_scheme:{scheme}"));
62        }
63        Ok(Self::ObjectStoreUri { scheme, uri })
64    }
65
66    /// Path or URI string passed to lower-level readers.
67    pub fn reader_path(&self) -> String {
68        match self {
69            Self::LocalPath(path) => path.display().to_string(),
70            Self::ObjectStoreUri { uri, .. } => uri.clone(),
71        }
72    }
73
74    /// Stable read-capability diagnostic string.
75    pub fn read_capability(&self) -> String {
76        match self {
77            Self::LocalPath(_) => "local_file_system".to_string(),
78            Self::ObjectStoreUri { scheme, .. } => format!("object_store:{scheme}"),
79        }
80    }
81
82    /// Whether this location is a local filesystem path.
83    pub fn is_local_path(&self) -> bool {
84        matches!(self, Self::LocalPath(_))
85    }
86
87    /// Borrow the local path when this location is local.
88    pub fn as_local_path(&self) -> Option<&Path> {
89        match self {
90            Self::LocalPath(path) => Some(path.as_path()),
91            Self::ObjectStoreUri { .. } => None,
92        }
93    }
94
95    /// Redacted diagnostic location. Local paths are preserved; object-store
96    /// URIs are represented by scheme plus stable hash only.
97    pub fn diagnostic_location(&self) -> String {
98        match self {
99            Self::LocalPath(path) => path.display().to_string(),
100            Self::ObjectStoreUri { scheme, uri } => object_store_diagnostic_id(scheme, uri),
101        }
102    }
103
104    #[doc(hidden)]
105    pub fn diagnostic_stable_id(&self, stable_id: &str) -> String {
106        match self {
107            Self::LocalPath(_) => redacted_stable_id_diagnostic(stable_id),
108            Self::ObjectStoreUri { scheme, .. } => {
109                format!("object_store:{scheme}:stable_id_fnv64:{:016x}", fnv1a64(stable_id.as_bytes()))
110            },
111        }
112    }
113}
114
115impl fmt::Display for ParquetSourceLocation {
116    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
117        match self {
118            Self::LocalPath(path) => write!(f, "{}", path.display()),
119            Self::ObjectStoreUri { uri, .. } => f.write_str(uri),
120        }
121    }
122}
123
124pub(super) fn parquet_source_location_for_original(location: &str) -> Result<ParquetSourceLocation, String> {
125    if let Some(path) = location.strip_prefix("file://localhost/") {
126        return Ok(ParquetSourceLocation::local_path(format!("/{path}")));
127    }
128    if let Some(path) = location.strip_prefix("file://") {
129        return Ok(ParquetSourceLocation::local_path(path));
130    }
131    if let Some(path) = location.strip_prefix("file:") {
132        return Ok(ParquetSourceLocation::local_path(path));
133    }
134    if location.contains("://") {
135        return ParquetSourceLocation::object_store_uri(location);
136    }
137    Ok(ParquetSourceLocation::local_path(location))
138}
139
140pub(super) fn object_store_reader_path(uri: &str) -> String {
141    uri.to_string()
142}
query_planning/source/location.rs

query_planning/source/
location.rs
