Skip to main content

query_planning/source/
read_policy.rs

1//! Parquet reader strategy, credentials, and cache-policy contracts.
2
3use std::fmt;
4
5use super::diagnostics::{redacted_debug_option, redacted_debug_value};
6
7/// High-level Parquet read strategy selected before source execution.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum ParquetReadStrategy {
10    /// Path-based cuDF/libcudf read for a local filesystem path.
11    LocalKvikio,
12    /// Read a previously materialized local byte-cache entry through the path-based reader.
13    LocalCacheKvikio,
14    /// Path-based cuDF/libcudf remote KvikIO read for an object-store URI such as `s3://...`.
15    RemoteDirect,
16}
17
18impl ParquetReadStrategy {
19    /// Stable diagnostic string.
20    pub const fn as_str(self) -> &'static str {
21        match self {
22            Self::LocalKvikio => "local_kvikio",
23            Self::LocalCacheKvikio => "local_cache_kvikio",
24            Self::RemoteDirect => "remote_direct",
25        }
26    }
27}
28
29impl fmt::Display for ParquetReadStrategy {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        f.write_str(self.as_str())
32    }
33}
34
35/// Tagged cuDF/KvikIO data-plane input.
36///
37/// Host-memory Parquet payloads are intentionally not representable.
38#[derive(Clone, PartialEq, Eq, Hash)]
39pub enum ResolvedParquetInput {
40    /// Local filesystem path, including a completed lakehouse-cache entry.
41    LocalKvikio {
42        /// Local path passed to the native reader, including a completed cache path.
43        reader_path: String,
44    },
45    /// Remote object-store URI read by cuDF's explicit KvikIO datasource.
46    RemoteKvikio {
47        /// Object-store URI passed to the native datasource.
48        uri: String,
49        /// Credential acquisition policy selected during source resolution.
50        credential_mode: CredentialMode,
51        /// Named storage profile, or `None` when no explicit profile was selected.
52        storage_profile_id: Option<String>,
53    },
54}
55
56impl fmt::Debug for ResolvedParquetInput {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        match self {
59            Self::LocalKvikio { reader_path } => f
60                .debug_struct("LocalKvikio")
61                .field("reader_path", &redacted_debug_value("reader_path", reader_path))
62                .finish(),
63            Self::RemoteKvikio { uri, credential_mode, storage_profile_id } => f
64                .debug_struct("RemoteKvikio")
65                .field("uri", &redacted_debug_value("uri", uri))
66                .field("credential_mode", credential_mode)
67                .field("storage_profile_id", storage_profile_id)
68                .finish(),
69        }
70    }
71}
72
73impl ResolvedParquetInput {
74    pub(crate) fn reader_location(&self) -> &str {
75        match self {
76            Self::LocalKvikio { reader_path } => reader_path,
77            Self::RemoteKvikio { uri, .. } => uri,
78        }
79    }
80}
81
82#[derive(Debug, Clone, PartialEq, Eq, Hash)]
83struct ObjectStoreAccess {
84    credential_mode: CredentialMode,
85    storage_profile_id: Option<String>,
86}
87
88/// Credential scope used by a resolved source read.
89#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
90pub enum CredentialMode {
91    /// No external credentials are needed.
92    #[default]
93    None,
94    /// Reader uses credentials already configured for the running process.
95    SharedProcessEnv,
96    /// Engine resolves per-profile credentials and scopes them to the reader.
97    ScopedResolved,
98}
99
100impl CredentialMode {
101    /// Stable diagnostic string.
102    pub const fn as_str(self) -> &'static str {
103        match self {
104            Self::None => "none",
105            Self::SharedProcessEnv => "shared_process_env",
106            Self::ScopedResolved => "scoped_resolved",
107        }
108    }
109}
110
111impl fmt::Display for CredentialMode {
112    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
113        f.write_str(self.as_str())
114    }
115}
116
117/// Local byte-cache status for a resolved Parquet source.
118#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
119pub enum LakehouseCacheStatus {
120    /// Cache strategy is not enabled for this source.
121    Disabled,
122    /// A complete cache entry is used as the reader input.
123    Hit,
124    /// A cache entry was absent.
125    Miss,
126    /// Cache fill was selected before reading.
127    FillBeforeRead,
128    /// Cache lookup was explicitly bypassed for benchmark isolation.
129    Bypassed,
130    /// Cache was considered but not used.
131    Skipped,
132}
133
134impl LakehouseCacheStatus {
135    /// Stable diagnostic string.
136    pub const fn as_str(self) -> &'static str {
137        match self {
138            Self::Disabled => "disabled",
139            Self::Hit => "hit",
140            Self::Miss => "miss",
141            Self::FillBeforeRead => "fill_before_read",
142            Self::Bypassed => "bypassed",
143            Self::Skipped => "skipped",
144        }
145    }
146}
147
148impl fmt::Display for LakehouseCacheStatus {
149    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150        f.write_str(self.as_str())
151    }
152}
153
154/// Foreground lakehouse cache read behavior selected for a resolved source.
155#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
156pub enum LakehouseCacheReadMode {
157    /// No complete-object cache lookup participates in the foreground read.
158    #[default]
159    Disabled,
160    /// Complete-object cache hits are eligible for foreground reads.
161    Enabled,
162    /// Complete-object cache hits were explicitly bypassed.
163    Bypassed,
164}
165
166impl LakehouseCacheReadMode {
167    /// Stable diagnostic string.
168    pub const fn as_str(self) -> &'static str {
169        match self {
170            Self::Disabled => "disabled",
171            Self::Enabled => "enabled",
172            Self::Bypassed => "bypassed",
173        }
174    }
175}
176
177impl fmt::Display for LakehouseCacheReadMode {
178    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
179        f.write_str(self.as_str())
180    }
181}
182
183/// Full-object lakehouse cache fill behavior selected for a resolved source.
184#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
185pub enum LakehouseCacheFillMode {
186    /// Do not fill a missing complete cache entry.
187    #[default]
188    Disabled,
189    /// Fill the complete object asynchronously after the foreground read is scheduled.
190    Background,
191    /// Fill the complete cache entry before opening it as a local reader path.
192    Blocking,
193}
194
195impl LakehouseCacheFillMode {
196    /// Stable diagnostic string.
197    pub const fn as_str(self) -> &'static str {
198        match self {
199            Self::Disabled => "disabled",
200            Self::Background => "background",
201            Self::Blocking => "blocking",
202        }
203    }
204}
205
206/// Resolved Parquet reader contract carried into native source execution.
207#[derive(Clone, PartialEq, Eq, Hash)]
208pub struct ResolvedParquetSource {
209    /// Frontend-supplied file location before read-strategy resolution.
210    pub original_location: String,
211    /// The only data-plane input supplied to cuDF.
212    pub input: ResolvedParquetInput,
213    /// Credential-free storage domain used for cache scoping, when known.
214    pub storage_domain: Option<String>,
215    /// Credential-free content identity used for cache scoping, when known.
216    pub content_identity: Option<String>,
217    /// Object-store identity retained only for cache fill and metadata maintenance
218    /// when the data-plane input is already a local path.
219    maintenance_access: Option<ObjectStoreAccess>,
220    /// Local byte-cache status for this source.
221    pub cache_status: LakehouseCacheStatus,
222    /// Foreground cache-read behavior for this source.
223    pub lakehouse_cache_read_mode: LakehouseCacheReadMode,
224    /// Lakehouse byte-cache root used by cache maintenance, when configured.
225    pub lakehouse_cache_root: Option<String>,
226    /// Complete-object cache path used for hit reads or background fills.
227    pub lakehouse_cache_path: Option<String>,
228    /// Best-effort retained byte target for the lakehouse byte-cache root.
229    pub lakehouse_cache_max_bytes: Option<u64>,
230    /// Cache-fill mode for complete-object lakehouse cache maintenance.
231    pub lakehouse_cache_fill_mode: LakehouseCacheFillMode,
232}
233
234impl fmt::Debug for ResolvedParquetSource {
235    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
236        f.debug_struct("ResolvedParquetSource")
237            .field("original_location", &redacted_debug_value("location", &self.original_location))
238            .field("reader_location", &redacted_debug_value("reader_location", self.reader_location()))
239            .field("storage_domain", &self.storage_domain)
240            .field("content_identity", &self.content_identity)
241            .field("storage_profile_id", &self.storage_profile_id())
242            .field("input", &self.input)
243            .field("read_strategy", &self.read_strategy())
244            .field("credential_mode", &self.credential_mode())
245            .field("cache_status", &self.cache_status)
246            .field("lakehouse_cache_read_mode", &self.lakehouse_cache_read_mode)
247            .field("lakehouse_cache_root", &redacted_debug_option("cache_root", self.lakehouse_cache_root.as_deref()))
248            .field("lakehouse_cache_path", &redacted_debug_option("cache_path", self.lakehouse_cache_path.as_deref()))
249            .field("lakehouse_cache_max_bytes", &self.lakehouse_cache_max_bytes)
250            .field("lakehouse_cache_fill_mode", &self.lakehouse_cache_fill_mode)
251            .finish()
252    }
253}
254
255impl ResolvedParquetSource {
256    /// Builds a local path-based read contract.
257    pub fn local_kvikio(location: impl Into<String>) -> Self {
258        let location = location.into();
259        Self {
260            original_location: location.clone(),
261            input: ResolvedParquetInput::LocalKvikio { reader_path: location },
262            storage_domain: None,
263            content_identity: None,
264            maintenance_access: None,
265            cache_status: LakehouseCacheStatus::Disabled,
266            lakehouse_cache_read_mode: LakehouseCacheReadMode::Disabled,
267            lakehouse_cache_root: None,
268            lakehouse_cache_path: None,
269            lakehouse_cache_max_bytes: None,
270            lakehouse_cache_fill_mode: LakehouseCacheFillMode::Disabled,
271        }
272    }
273
274    /// Builds a read contract for a local byte-cache entry, with the given cache status.
275    pub fn local_cache_kvikio(
276        original_location: impl Into<String>,
277        reader_location: impl Into<String>,
278        credential_mode: CredentialMode,
279        cache_status: LakehouseCacheStatus,
280    ) -> Self {
281        Self {
282            original_location: original_location.into(),
283            input: ResolvedParquetInput::LocalKvikio { reader_path: reader_location.into() },
284            storage_domain: None,
285            content_identity: None,
286            maintenance_access: Some(ObjectStoreAccess { credential_mode, storage_profile_id: None }),
287            cache_status,
288            lakehouse_cache_read_mode: LakehouseCacheReadMode::Enabled,
289            lakehouse_cache_root: None,
290            lakehouse_cache_path: None,
291            lakehouse_cache_max_bytes: None,
292            lakehouse_cache_fill_mode: LakehouseCacheFillMode::Disabled,
293        }
294    }
295
296    /// Builds a remote path-based read contract where cuDF/libcudf owns the
297    /// object-store read. Host-memory Parquet payloads are not representable.
298    pub fn remote_kvikio(original_location: impl Into<String>, credential_mode: CredentialMode) -> Self {
299        let original_location = original_location.into();
300        Self {
301            input: ResolvedParquetInput::RemoteKvikio {
302                uri: original_location.clone(),
303                credential_mode,
304                storage_profile_id: None,
305            },
306            original_location,
307            storage_domain: None,
308            content_identity: None,
309            maintenance_access: None,
310            cache_status: LakehouseCacheStatus::Disabled,
311            lakehouse_cache_read_mode: LakehouseCacheReadMode::Disabled,
312            lakehouse_cache_root: None,
313            lakehouse_cache_path: None,
314            lakehouse_cache_max_bytes: None,
315            lakehouse_cache_fill_mode: LakehouseCacheFillMode::Disabled,
316        }
317    }
318
319    /// Replaces the local byte-cache status after strategy fallback.
320    pub fn with_cache_status(mut self, cache_status: LakehouseCacheStatus) -> Self {
321        self.cache_status = cache_status;
322        self
323    }
324
325    /// Selects foreground complete-object lakehouse cache read behavior.
326    pub fn with_lakehouse_cache_read_mode(mut self, read_mode: LakehouseCacheReadMode) -> Self {
327        self.lakehouse_cache_read_mode = read_mode;
328        self
329    }
330
331    /// Attaches a credential-free storage-domain identity for cache scoping.
332    pub fn with_storage_domain(mut self, storage_domain: impl Into<String>) -> Self {
333        self.storage_domain = Some(storage_domain.into());
334        self
335    }
336
337    /// Attaches a credential-free content identity for cache scoping.
338    pub fn with_content_identity(mut self, content_identity: impl Into<String>) -> Self {
339        self.content_identity = Some(content_identity.into());
340        self
341    }
342
343    /// Attaches the execution-time source storage profile id.
344    pub fn with_storage_profile_id(mut self, profile_id: impl Into<String>) -> Self {
345        self.set_storage_profile_id(Some(profile_id.into()));
346        self
347    }
348
349    /// Attaches an execution-time source storage profile id when present.
350    pub fn with_optional_storage_profile_id(mut self, profile_id: Option<String>) -> Self {
351        self.set_storage_profile_id(profile_id);
352        self
353    }
354
355    fn set_storage_profile_id(&mut self, profile_id: Option<String>) {
356        match &mut self.input {
357            ResolvedParquetInput::RemoteKvikio { storage_profile_id, .. } => *storage_profile_id = profile_id,
358            ResolvedParquetInput::LocalKvikio { .. } => {
359                self.maintenance_access
360                    .get_or_insert(ObjectStoreAccess {
361                        credential_mode: CredentialMode::None,
362                        storage_profile_id: None,
363                    })
364                    .storage_profile_id = profile_id;
365            },
366        }
367    }
368
369    /// Attaches lakehouse byte-cache maintenance settings.
370    pub fn with_lakehouse_cache_config(
371        mut self,
372        cache_root: impl Into<String>,
373        max_retained_bytes: Option<u64>,
374    ) -> Self {
375        self.lakehouse_cache_root = Some(cache_root.into());
376        self.lakehouse_cache_max_bytes = max_retained_bytes;
377        self
378    }
379
380    /// Attaches a complete-object lakehouse cache path.
381    pub fn with_lakehouse_cache_path(mut self, cache_path: impl Into<String>) -> Self {
382        self.lakehouse_cache_path = Some(cache_path.into());
383        self
384    }
385
386    /// Selects complete-object lakehouse cache fill behavior.
387    pub fn with_lakehouse_cache_fill_mode(mut self, fill_mode: LakehouseCacheFillMode) -> Self {
388        self.lakehouse_cache_fill_mode = fill_mode;
389        self
390    }
391
392    /// Path or URI passed to cuDF.
393    pub fn reader_location(&self) -> &str {
394        self.input.reader_location()
395    }
396
397    /// Credential scope for the active remote read or local-cache maintenance.
398    pub fn credential_mode(&self) -> CredentialMode {
399        match &self.input {
400            ResolvedParquetInput::RemoteKvikio { credential_mode, .. } => *credential_mode,
401            ResolvedParquetInput::LocalKvikio { .. } => {
402                self.maintenance_access.as_ref().map_or(CredentialMode::None, |access| access.credential_mode)
403            },
404        }
405    }
406
407    /// Storage profile for the active remote read or local-cache maintenance.
408    pub fn storage_profile_id(&self) -> Option<&str> {
409        match &self.input {
410            ResolvedParquetInput::RemoteKvikio { storage_profile_id, .. } => storage_profile_id.as_deref(),
411            ResolvedParquetInput::LocalKvikio { .. } => {
412                self.maintenance_access.as_ref().and_then(|access| access.storage_profile_id.as_deref())
413            },
414        }
415    }
416
417    /// Report-facing strategy derived from the tagged input and cache facts.
418    pub fn read_strategy(&self) -> ParquetReadStrategy {
419        match self.input {
420            ResolvedParquetInput::RemoteKvikio { .. } => ParquetReadStrategy::RemoteDirect,
421            ResolvedParquetInput::LocalKvikio { .. }
422                if matches!(self.cache_status, LakehouseCacheStatus::Hit | LakehouseCacheStatus::FillBeforeRead) =>
423            {
424                ParquetReadStrategy::LocalCacheKvikio
425            },
426            ResolvedParquetInput::LocalKvikio { .. } => ParquetReadStrategy::LocalKvikio,
427        }
428    }
429}