polars_io/utils/
byte_source.rs

1use std::ops::Range;
2use std::path::Path;
3use std::sync::Arc;
4
5use polars_buffer::Buffer;
6use polars_core::prelude::PlHashMap;
7use polars_error::{PolarsResult, feature_gated};
8use polars_utils::_limit_path_len_io_err;
9use polars_utils::mmap::MMapSemaphore;
10use polars_utils::pl_path::PlRefPath;
11
12use crate::cloud::concurrency_config::{ConcurrencyStrategy, FetchConfig};
13use crate::cloud::options::CloudOptions;
14#[cfg(feature = "cloud")]
15use crate::cloud::{
16    CloudLocation, ObjectStorePath, PolarsObjectStore, build_object_store, object_path_from_str,
17};
18use crate::metrics::IOMetrics;
19
20#[allow(async_fn_in_trait)]
21pub trait ByteSource: Send + Sync {
22    async fn get_size(&self) -> PolarsResult<usize>;
23    /// # Panics
24    /// Panics if `range` is not in bounds.
25    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>>;
26    /// Note: This will mutably sort ranges for coalescing.
27    async fn get_ranges(
28        &self,
29        ranges: &mut [Range<usize>],
30    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>>;
31}
32
33/// Byte source backed by a `Buffer`, which can potentially be memory-mapped.
34pub struct BufferByteSource(pub Buffer<u8>);
35
36impl BufferByteSource {
37    async fn try_new_mmap_from_path(
38        path: &Path,
39        _cloud_options: Option<&CloudOptions>,
40    ) -> PolarsResult<Self> {
41        let file = Arc::new(
42            tokio::fs::File::open(path)
43                .await
44                .map_err(|err| _limit_path_len_io_err(path, err))?
45                .into_std()
46                .await,
47        );
48
49        Ok(Self(Buffer::from_owner(MMapSemaphore::new_from_file(
50            &file,
51        )?)))
52    }
53}
54
55impl ByteSource for BufferByteSource {
56    async fn get_size(&self) -> PolarsResult<usize> {
57        Ok(self.0.as_ref().len())
58    }
59
60    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
61        let out = self.0.clone().sliced(range);
62        Ok(out)
63    }
64
65    async fn get_ranges(
66        &self,
67        ranges: &mut [Range<usize>],
68    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
69        Ok(ranges
70            .iter()
71            .map(|x| (x.start, self.0.clone().sliced(x.clone())))
72            .collect())
73    }
74}
75
76#[cfg(feature = "cloud")]
77pub struct ObjectStoreByteSource {
78    store: PolarsObjectStore,
79    path: ObjectStorePath,
80    config: FetchConfig,
81}
82
83#[cfg(feature = "cloud")]
84impl ObjectStoreByteSource {
85    async fn try_new_from_path(
86        path: PlRefPath,
87        cloud_options: Option<&CloudOptions>,
88        io_metrics: Option<Arc<IOMetrics>>,
89        config: FetchConfig,
90    ) -> PolarsResult<Self> {
91        let (CloudLocation { prefix, .. }, mut store) =
92            build_object_store(path, cloud_options, false).await?;
93        let path = object_path_from_str(&prefix)?;
94
95        store.set_io_metrics(io_metrics);
96
97        Ok(Self {
98            store,
99            path,
100            config,
101        })
102    }
103
104    #[allow(unused)]
105    fn chunk_size(&self) -> usize {
106        self.config.chunk_size
107    }
108
109    fn concurrency_strategy(&self) -> ConcurrencyStrategy {
110        self.config.strategy
111    }
112}
113
114#[cfg(feature = "cloud")]
115impl ByteSource for ObjectStoreByteSource {
116    async fn get_size(&self) -> PolarsResult<usize> {
117        Ok(self
118            .store
119            .head(&self.path, ConcurrencyStrategy::Legacy)
120            .await?
121            .size as usize)
122    }
123
124    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
125        self.store.get_range(&self.path, range, self.config).await
126    }
127
128    async fn get_ranges(
129        &self,
130        ranges: &mut [Range<usize>],
131    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
132        self.store
133            .get_ranges_sort(&self.path, ranges, self.config)
134            .await
135    }
136}
137
138/// Dynamic dispatch to async functions.
139pub enum DynByteSource {
140    Buffer(BufferByteSource),
141    #[cfg(feature = "cloud")]
142    Cloud(ObjectStoreByteSource),
143}
144
145impl DynByteSource {
146    pub fn variant_name(&self) -> &str {
147        match self {
148            Self::Buffer(_) => "Buffer",
149            #[cfg(feature = "cloud")]
150            Self::Cloud(_) => "Cloud",
151        }
152    }
153
154    pub fn is_cloud(&self) -> bool {
155        match self {
156            Self::Buffer(_) => false,
157            #[cfg(feature = "cloud")]
158            Self::Cloud(_) => true,
159        }
160    }
161
162    pub fn chunk_size(&self) -> Option<usize> {
163        match self {
164            Self::Buffer(_) => None,
165            #[cfg(feature = "cloud")]
166            Self::Cloud(source) => Some(source.config.chunk_size),
167        }
168    }
169
170    pub fn concurrency_strategy(&self) -> Option<ConcurrencyStrategy> {
171        match self {
172            Self::Buffer(_) => None,
173            #[cfg(feature = "cloud")]
174            Self::Cloud(source) => Some(source.concurrency_strategy()),
175        }
176    }
177}
178
179impl Default for DynByteSource {
180    fn default() -> Self {
181        Self::Buffer(BufferByteSource(Buffer::new()))
182    }
183}
184
185impl ByteSource for DynByteSource {
186    async fn get_size(&self) -> PolarsResult<usize> {
187        match self {
188            Self::Buffer(v) => v.get_size().await,
189            #[cfg(feature = "cloud")]
190            Self::Cloud(v) => v.get_size().await,
191        }
192    }
193
194    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
195        match self {
196            Self::Buffer(v) => v.get_range(range).await,
197            #[cfg(feature = "cloud")]
198            Self::Cloud(v) => v.get_range(range).await,
199        }
200    }
201
202    async fn get_ranges(
203        &self,
204        ranges: &mut [Range<usize>],
205    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
206        match self {
207            Self::Buffer(v) => v.get_ranges(ranges).await,
208            #[cfg(feature = "cloud")]
209            Self::Cloud(v) => v.get_ranges(ranges).await,
210        }
211    }
212}
213
214impl From<BufferByteSource> for DynByteSource {
215    fn from(value: BufferByteSource) -> Self {
216        Self::Buffer(value)
217    }
218}
219
220#[cfg(feature = "cloud")]
221impl From<ObjectStoreByteSource> for DynByteSource {
222    fn from(value: ObjectStoreByteSource) -> Self {
223        Self::Cloud(value)
224    }
225}
226
227impl From<Buffer<u8>> for DynByteSource {
228    fn from(value: Buffer<u8>) -> Self {
229        Self::Buffer(BufferByteSource(value))
230    }
231}
232
233#[derive(Clone, Debug)]
234pub enum DynByteSourceBuilder {
235    Mmap,
236    /// Supports both cloud and local files, requires cloud feature.
237    ObjectStore(FetchConfig),
238}
239
240impl DynByteSourceBuilder {
241    pub async fn try_build_from_path(
242        &self,
243        path: PlRefPath,
244        cloud_options: Option<&CloudOptions>,
245        io_metrics: Option<Arc<IOMetrics>>,
246    ) -> PolarsResult<DynByteSource> {
247        Ok(match *self {
248            Self::Mmap => {
249                BufferByteSource::try_new_mmap_from_path(path.as_std_path(), cloud_options)
250                    .await?
251                    .into()
252            },
253            Self::ObjectStore(fetch_config) => feature_gated!("cloud", {
254                ObjectStoreByteSource::try_new_from_path(
255                    path,
256                    cloud_options,
257                    io_metrics,
258                    fetch_config,
259                )
260                .await?
261                .into()
262            }),
263        })
264    }
265
266    pub fn chunk_size(&self) -> Option<usize> {
267        match self {
268            Self::Mmap => None,
269            Self::ObjectStore(fetch_config) => Some(fetch_config.chunk_size),
270        }
271    }
272
273    pub fn concurrency_strategy(&self) -> Option<&ConcurrencyStrategy> {
274        match self {
275            Self::Mmap => None,
276            Self::ObjectStore(fetch_config) => Some(&fetch_config.strategy),
277        }
278    }
279}
polars_io/utils/byte_source.rs

polars_io/utils/
byte_source.rs
