polars_io/file_cache/
utils.rs

1use std::sync::{Arc, LazyLock};
2use std::time::UNIX_EPOCH;
3
4use polars_error::{PolarsError, PolarsResult};
5use polars_utils::aliases::PlHashMap;
6use polars_utils::pl_path::{CloudScheme, PlRefPath};
7
8use super::cache::{FILE_CACHE, get_env_file_cache_ttl};
9use super::entry::FileCacheEntry;
10use super::file_fetcher::{CloudFileFetcher, LocalFileFetcher};
11use crate::cloud::{
12    CloudLocation, CloudOptions, PolarsObjectStore, build_object_store, object_path_from_str,
13};
14use crate::path_utils::{POLARS_TEMP_DIR_BASE_PATH, ensure_directory_init};
15
16pub static FILE_CACHE_PREFIX: LazyLock<PlRefPath> = LazyLock::new(|| {
17    let path = PlRefPath::try_from_path(&POLARS_TEMP_DIR_BASE_PATH.join("file-cache/")).unwrap();
18
19    if let Err(err) = ensure_directory_init(path.as_ref()) {
20        panic!(
21            "failed to create file cache directory: path = {}, err = {}",
22            path, err
23        );
24    }
25
26    path
27});
28
29pub(super) fn last_modified_u64(metadata: &std::fs::Metadata) -> u64 {
30    u64::try_from(
31        metadata
32            .modified()
33            .unwrap()
34            .duration_since(UNIX_EPOCH)
35            .unwrap()
36            .as_millis(),
37    )
38    .unwrap()
39}
40
41pub(super) fn update_last_accessed(file: &std::fs::File) {
42    let file_metadata = file.metadata().unwrap();
43
44    if let Err(e) = file.set_times(
45        std::fs::FileTimes::new()
46            .set_modified(file_metadata.modified().unwrap())
47            .set_accessed(std::time::SystemTime::now()),
48    ) {
49        panic!("failed to update file last accessed time: {e}");
50    }
51}
52
53pub async fn init_entries_from_uri_list(
54    uri_list: impl ExactSizeIterator<Item = PlRefPath> + Send + 'static,
55    cloud_options: Option<&CloudOptions>,
56) -> PolarsResult<Vec<Arc<FileCacheEntry>>> {
57    init_entries_from_uri_list_impl(Box::new(uri_list), cloud_options).await
58}
59
60async fn init_entries_from_uri_list_impl(
61    uri_list: Box<dyn ExactSizeIterator<Item = PlRefPath> + Send + 'static>,
62    cloud_options: Option<&CloudOptions>,
63) -> PolarsResult<Vec<Arc<FileCacheEntry>>> {
64    #[allow(clippy::len_zero)]
65    if uri_list.len() == 0 {
66        return Ok(Default::default());
67    }
68
69    let mut uri_list = uri_list.peekable();
70
71    let first_uri = uri_list.peek().unwrap().clone();
72
73    let file_cache_ttl = cloud_options
74        .map(|x| x.file_cache_ttl)
75        .unwrap_or_else(get_env_file_cache_ttl);
76
77    if first_uri.has_scheme() {
78        let uri_list: Vec<PlRefPath> = uri_list.collect();
79
80        // Http URIs can differ in origin.
81        let is_http = matches!(
82            first_uri.scheme(),
83            Some(CloudScheme::Http | CloudScheme::Https)
84        );
85
86        let authorities: Vec<&str> = if is_http {
87            Vec::new()
88        } else {
89            uri_list
90                .iter()
91                .map(|uri| &uri.as_str()[..uri.authority_end_position()])
92                .collect()
93        };
94
95        // One object store per bucket, held here so that global cache evictions cannot
96        // affect this call.
97        let mut representatives: PlHashMap<&str, &PlRefPath> = PlHashMap::default();
98        for (uri, authority) in uri_list.iter().zip(&authorities) {
99            representatives.entry(authority).or_insert(uri);
100        }
101
102        let shared_object_stores: PlHashMap<&str, PolarsObjectStore> =
103            futures::future::try_join_all(representatives.into_iter().map(
104                |(authority, uri)| async move {
105                    let (_, object_store) =
106                        build_object_store(uri.clone(), cloud_options, false).await?;
107                    PolarsResult::Ok((authority, object_store))
108                },
109            ))
110            .await?
111            .into_iter()
112            .collect();
113
114        futures::future::try_join_all(uri_list.iter().enumerate().map(|(i, uri)| {
115            let uri = uri.clone();
116            let shared_object_store = authorities
117                .get(i)
118                .and_then(|authority| shared_object_stores.get(authority))
119                .cloned();
120
121            async move {
122                let object_store = if let Some(shared_object_store) = shared_object_store {
123                    shared_object_store
124                } else {
125                    let (_, object_store) =
126                        build_object_store(uri.clone(), cloud_options, false).await?;
127                    object_store
128                };
129
130                FILE_CACHE.init_entry(
131                    uri.clone(),
132                    &|| {
133                        let CloudLocation { prefix, .. } =
134                            CloudLocation::new(uri.clone(), false).unwrap();
135                        let cloud_path = object_path_from_str(&prefix)?;
136                        let object_store = object_store.clone();
137
138                        Ok(Arc::new(CloudFileFetcher {
139                            uri: uri.clone(),
140                            object_store,
141                            cloud_path,
142                        }))
143                    },
144                    file_cache_ttl,
145                )
146            }
147        }))
148        .await
149    } else {
150        let mut out = Vec::with_capacity(uri_list.len());
151        for uri in uri_list {
152            let uri = tokio::fs::canonicalize(uri.as_str()).await.map_err(|err| {
153                let msg = Some(format!("{}: {}", err, uri).into());
154                PolarsError::IO {
155                    error: err.into(),
156                    msg,
157                }
158            })?;
159            let uri = PlRefPath::try_from_pathbuf(uri)?;
160
161            out.push(FILE_CACHE.init_entry(
162                uri.clone(),
163                &|| Ok(Arc::new(LocalFileFetcher::from_uri(uri.clone()))),
164                file_cache_ttl,
165            )?)
166        }
167        Ok(out)
168    }
169}
polars_io/file_cache/utils.rs

polars_io/file_cache/
utils.rs
