1use std::fmt;
4
5use super::diagnostics::{redacted_debug_option, redacted_debug_value};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum ParquetReadStrategy {
10 LocalKvikio,
12 LocalCacheKvikio,
14 RemoteDirect,
16}
17
18impl ParquetReadStrategy {
19 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#[derive(Clone, PartialEq, Eq, Hash)]
39pub enum ResolvedParquetInput {
40 LocalKvikio {
42 reader_path: String,
44 },
45 RemoteKvikio {
47 uri: String,
49 credential_mode: CredentialMode,
51 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#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
90pub enum CredentialMode {
91 #[default]
93 None,
94 SharedProcessEnv,
96 ScopedResolved,
98}
99
100impl CredentialMode {
101 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
119pub enum LakehouseCacheStatus {
120 Disabled,
122 Hit,
124 Miss,
126 FillBeforeRead,
128 Bypassed,
130 Skipped,
132}
133
134impl LakehouseCacheStatus {
135 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#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
156pub enum LakehouseCacheReadMode {
157 #[default]
159 Disabled,
160 Enabled,
162 Bypassed,
164}
165
166impl LakehouseCacheReadMode {
167 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#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
185pub enum LakehouseCacheFillMode {
186 #[default]
188 Disabled,
189 Background,
191 Blocking,
193}
194
195impl LakehouseCacheFillMode {
196 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#[derive(Clone, PartialEq, Eq, Hash)]
208pub struct ResolvedParquetSource {
209 pub original_location: String,
211 pub input: ResolvedParquetInput,
213 pub storage_domain: Option<String>,
215 pub content_identity: Option<String>,
217 maintenance_access: Option<ObjectStoreAccess>,
220 pub cache_status: LakehouseCacheStatus,
222 pub lakehouse_cache_read_mode: LakehouseCacheReadMode,
224 pub lakehouse_cache_root: Option<String>,
226 pub lakehouse_cache_path: Option<String>,
228 pub lakehouse_cache_max_bytes: Option<u64>,
230 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 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 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 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 pub fn with_cache_status(mut self, cache_status: LakehouseCacheStatus) -> Self {
321 self.cache_status = cache_status;
322 self
323 }
324
325 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 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 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 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 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 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 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 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 pub fn reader_location(&self) -> &str {
394 self.input.reader_location()
395 }
396
397 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 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 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}