polars_io/utils/
byte_source.rs

1use std::ops::Range;
2#[cfg(target_os = "linux")]
3use std::os::fd::AsRawFd;
4use std::path::Path;
5use std::sync::{Arc, LazyLock};
6
7use dio_align::DioAlign;
8use futures::{StreamExt, TryStreamExt};
9use polars_buffer::Buffer;
10use polars_config::FileAdvice;
11use polars_core::prelude::PlHashMap;
12use polars_core::runtime::ASYNC;
13use polars_error::{PolarsResult, feature_gated, polars_bail, polars_err};
14use polars_utils::aliases::InitHashMaps;
15use polars_utils::io::_limit_path_len_io_err;
16use polars_utils::mmap::MMapSemaphore;
17use polars_utils::pl_path::PlRefPath;
18use tokio::sync::Semaphore;
19
20use crate::cloud::concurrency_config::{ConcurrencyStrategy, FetchConfig};
21use crate::cloud::options::CloudOptions;
22#[cfg(feature = "cloud")]
23use crate::cloud::{
24    CloudLocation, ObjectStorePath, PolarsObjectStore, build_object_store, object_path_from_str,
25};
26use crate::metrics::{IOMetrics, OptIOMetrics};
27
28pub mod dio_align;
29#[cfg(target_os = "linux")]
30mod direct_io;
31#[allow(async_fn_in_trait)]
32pub trait ByteSource: Send + Sync {
33    async fn get_size(&self) -> PolarsResult<usize>;
34    /// Fetch the last `n` bytes and the total size of the source, in a single request. Returns
35    /// fewer than `n` bytes if the source is smaller than `n`.
36    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)>;
37    /// # Panics
38    /// Panics if `range` is not in bounds.
39    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>>;
40    /// Note: This will mutably sort ranges for coalescing.
41    async fn get_ranges(
42        &self,
43        ranges: &mut [Range<usize>],
44    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>>;
45}
46
47/// Byte source backed by a `Buffer`, which can potentially be memory-mapped.
48pub struct BufferByteSource(pub Buffer<u8>);
49
50impl BufferByteSource {
51    async fn try_new_mmap_from_path(
52        path: &Path,
53        _cloud_options: Option<&CloudOptions>,
54    ) -> PolarsResult<Self> {
55        let file = Arc::new(
56            tokio::fs::File::open(path)
57                .await
58                .map_err(|err| _limit_path_len_io_err(path, err))?
59                .into_std()
60                .await,
61        );
62
63        Ok(Self(Buffer::from_owner(MMapSemaphore::new_from_file(
64            &file,
65        )?)))
66    }
67}
68
69impl ByteSource for BufferByteSource {
70    async fn get_size(&self) -> PolarsResult<usize> {
71        Ok(self.0.as_ref().len())
72    }
73
74    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
75        let len = self.0.as_ref().len();
76        Ok((self.0.clone().sliced(len.saturating_sub(n)..len), len))
77    }
78
79    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
80        let out = self.0.clone().sliced(range);
81        Ok(out)
82    }
83
84    async fn get_ranges(
85        &self,
86        ranges: &mut [Range<usize>],
87    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
88        Ok(ranges
89            .iter()
90            .map(|x| (x.start, self.0.clone().sliced(x.clone())))
91            .collect())
92    }
93}
94
95/// Byte source backed by a `File`.
96pub struct FileByteSource {
97    file: Arc<std::fs::File>,
98    // Alignment for O_DIRECT, if supported.
99    o_direct_align: Option<DioAlign>,
100    // Manage concurrency.
101    concurrency: usize,
102    permits: Arc<Semaphore>,
103    // File size.
104    size: u64,
105    io_metrics: OptIOMetrics,
106}
107
108/// Each permit pins a tokio blocking thread for the duration of a `pread`.
109pub fn global_read_permits() -> Arc<Semaphore> {
110    static PERMITS: LazyLock<Arc<Semaphore>> = LazyLock::new(|| {
111        Arc::new(Semaphore::new(
112            polars_config::config().file_read_concurrency().max(1) as usize,
113        ))
114    });
115
116    PERMITS.clone()
117}
118
119#[derive(Clone)]
120pub struct FileReadContext {
121    pub enable_o_direct: bool,
122    pub concurrency: usize,
123    pub permits: Arc<Semaphore>,
124    /// Ignored when direct I/O is active: there is no page cache to advise.
125    pub advice: FileAdvice,
126}
127
128impl std::fmt::Debug for FileReadContext {
129    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
130        f.debug_struct("FileReadContext")
131            .field("available_permits", &self.permits.available_permits())
132            .field("enable_odirect", &self.enable_o_direct)
133            .field("concurrency", &self.concurrency)
134            .field("advice", &self.advice)
135            .finish()
136    }
137}
138
139#[cfg(target_os = "linux")]
140fn _fadvise(file: &std::fs::File, advice: FileAdvice) {
141    let advice = match advice {
142        // The OS default is already POSIX_FADV_NORMAL.
143        FileAdvice::Normal => return,
144        FileAdvice::Sequential => libc::POSIX_FADV_SEQUENTIAL,
145        FileAdvice::Random => libc::POSIX_FADV_RANDOM,
146        FileAdvice::WillNeed => libc::POSIX_FADV_WILLNEED,
147    };
148
149    unsafe {
150        libc::posix_fadvise(file.as_raw_fd(), 0, 0, advice);
151    }
152}
153
154#[cfg(all(unix, not(feature = "nightly")))]
155fn pread_exact(file: &std::fs::File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
156    use std::os::unix::fs::FileExt;
157    file.read_exact_at(buf, offset)
158}
159
160#[cfg(all(windows, not(feature = "nightly")))]
161fn pread_exact(file: &std::fs::File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
162    use std::os::windows::fs::FileExt;
163    let mut filled = 0;
164    while filled < buf.len() {
165        match file.seek_read(&mut buf[filled..], offset + filled as u64)? {
166            0 => return Err(std::io::ErrorKind::UnexpectedEof.into()),
167            n => filled += n,
168        }
169    }
170    Ok(())
171}
172
173/// `pread_exact` into uninitialized memory, so that `read_buffered` does not
174/// have to zero the buffer the read is about to overwrite anyway.
175#[cfg(all(feature = "nightly", unix))]
176fn pread_exact_uninit(
177    file: &std::fs::File,
178    buf: std::io::BorrowedCursor<'_, u8>,
179    offset: u64,
180) -> std::io::Result<()> {
181    use std::os::unix::fs::FileExt;
182    file.read_buf_exact_at(buf, offset)
183}
184
185#[cfg(all(feature = "nightly", windows))]
186fn pread_exact_uninit(
187    file: &std::fs::File,
188    mut buf: std::io::BorrowedCursor<'_, u8>,
189    mut offset: u64,
190) -> std::io::Result<()> {
191    use std::os::windows::fs::FileExt;
192    // `seek_read_buf` is the short-read variant; loop it to fill the cursor.
193    while buf.capacity() > 0 {
194        let written = buf.written();
195        file.seek_read_buf(buf.reborrow(), offset)?;
196        match buf.written() - written {
197            0 => return Err(std::io::ErrorKind::UnexpectedEof.into()),
198            n => offset += n as u64,
199        }
200    }
201    Ok(())
202}
203
204/// Read `len` bytes at `offset` through the page cache.
205///
206/// Under `nightly` the destination is left uninitialized and filled by the read
207/// itself; on stable it must be zeroed first, which memsets the whole range
208/// before the kernel overwrites it.
209#[cfg(feature = "nightly")]
210fn read_buffered(file: &std::fs::File, offset: u64, len: usize) -> PolarsResult<Buffer<u8>> {
211    let mut v: Vec<u8> = Vec::with_capacity(len);
212
213    let filled = {
214        let mut buf = std::io::BorrowedBuf::from(&mut v.spare_capacity_mut()[..len]);
215        pread_exact_uninit(file, buf.unfilled(), offset)?;
216        buf.len()
217    };
218
219    debug_assert_eq!(filled, len);
220    // Safety: `pread_exact_uninit` filled the cursor, so the first `filled`
221    // bytes of the allocation are initialized.
222    unsafe { v.set_len(filled) };
223
224    Ok(Buffer::from(v))
225}
226
227#[cfg(not(feature = "nightly"))]
228fn read_buffered(file: &std::fs::File, offset: u64, len: usize) -> PolarsResult<Buffer<u8>> {
229    let mut buf = vec![0u8; len];
230    pread_exact(file, &mut buf, offset)?;
231    Ok(Buffer::from(buf))
232}
233
234/// Read `len` bytes at `offset`, through direct I/O when `align` says the file
235/// supports it and through the page cache otherwise.
236#[cfg(target_os = "linux")]
237fn read_at(
238    file: &std::fs::File,
239    align: Option<DioAlign>,
240    offset: u64,
241    len: usize,
242    size: u64,
243) -> PolarsResult<Buffer<u8>> {
244    match align {
245        Some(align) => direct_io::read_aligned(file, align, offset, len, size),
246        None => read_buffered(file, offset, len),
247    }
248}
249
250#[cfg(not(target_os = "linux"))]
251fn read_at(
252    file: &std::fs::File,
253    _align: Option<DioAlign>,
254    offset: u64,
255    len: usize,
256    _size: u64,
257) -> PolarsResult<Buffer<u8>> {
258    read_buffered(file, offset, len)
259}
260
261impl FileByteSource {
262    async fn try_new_from_path(
263        path: PlRefPath,
264        read_context: FileReadContext,
265        io_metrics: Option<Arc<IOMetrics>>,
266    ) -> PolarsResult<Self> {
267        // The open path is `open`, `fcntl`, `statx` and `fstat` - all blocking.
268        ASYNC
269            .spawn_blocking(move || Self::open_blocking(path, read_context, io_metrics))
270            .await
271            .expect("blocking task panicked")
272    }
273
274    fn open_blocking(
275        path: PlRefPath,
276        read_context: FileReadContext,
277        io_metrics: Option<Arc<IOMetrics>>,
278    ) -> PolarsResult<Self> {
279        let path = path.as_std_path();
280        let enable_o_direct = read_context.enable_o_direct;
281
282        let file = {
283            #[cfg(target_os = "linux")]
284            let f = if enable_o_direct {
285                direct_io::open_o_direct(path).or_else(|e| match e.raw_os_error() {
286                    // The filesystem does not support O_DIRECT.
287                    Some(libc::EINVAL) => std::fs::File::open(path),
288                    _ => Err(e),
289                })
290            } else {
291                std::fs::File::open(path)
292            };
293            #[cfg(not(target_os = "linux"))]
294            let f = std::fs::File::open(path);
295
296            f.map_err(|e| _limit_path_len_io_err(path, e))?
297        };
298
299        let file = Arc::new(file);
300
301        #[cfg(target_os = "linux")]
302        let o_direct_align =
303            direct_io::probe_or_disable(&file).map_err(|e| _limit_path_len_io_err(path, e))?;
304
305        #[cfg(not(target_os = "linux"))]
306        let o_direct_align: Option<DioAlign> = None;
307
308        #[cfg(target_os = "linux")]
309        if o_direct_align.is_none() {
310            // Inert under O_DIRECT: there is no page cache to advise.
311            _fadvise(&file, read_context.advice);
312        }
313
314        let concurrency = read_context.concurrency;
315        let permits = read_context.permits;
316
317        let size = file
318            .metadata()
319            .map_err(|e| _limit_path_len_io_err(path, e))?
320            .len();
321
322        if polars_config::config().verbose() {
323            let name = if polars_config::config().verbose_sensitive() {
324                path.display().to_string()
325            } else {
326                "<file>".to_string()
327            };
328            match (enable_o_direct, o_direct_align) {
329                (true, Some(a)) => eprintln!(
330                    "[FileByteSource]: {name}: direct IO active, \
331                        alignment offset: {}, memory: {}",
332                    a.offset, a.memory
333                ),
334                (true, None) => eprintln!(
335                    "[FileByteSource]: {name}: direct IO requested but not active, \
336                        using buffered reads"
337                ),
338                (false, _) => {},
339            }
340        }
341
342        Ok(FileByteSource {
343            file,
344            o_direct_align,
345            concurrency,
346            permits,
347            size,
348            io_metrics: OptIOMetrics(io_metrics),
349        })
350    }
351
352    pub fn try_new_from_std(
353        file: std::fs::File,
354        read_context: FileReadContext,
355        io_metrics: Option<Arc<IOMetrics>>,
356    ) -> PolarsResult<Self> {
357        let size = file.metadata()?.len();
358
359        #[cfg(target_os = "linux")]
360        let o_direct_align = direct_io::probe_or_disable(&file)?;
361
362        #[cfg(not(target_os = "linux"))]
363        let o_direct_align: Option<DioAlign> = None;
364
365        #[cfg(target_os = "linux")]
366        if o_direct_align.is_none() {
367            _fadvise(&file, read_context.advice);
368        }
369
370        let concurrency = read_context.concurrency;
371        let permits = read_context.permits;
372
373        // For verbose logging only. We cannot enable since the file_handle was handed to us.
374        let enable_o_direct = read_context.enable_o_direct;
375
376        if polars_config::config().verbose() {
377            match (enable_o_direct, o_direct_align) {
378                (true, Some(a)) => eprintln!(
379                    "[FileByteSource]: direct IO active, alignment offset: {}, memory: {}",
380                    a.offset, a.memory
381                ),
382                (true, None) => eprintln!(
383                    "[FileByteSource]: direct IO requested but not active, using buffered reads",
384                ),
385                (false, _) => {},
386            }
387        }
388
389        Ok(Self {
390            file: Arc::new(file),
391            o_direct_align,
392            concurrency,
393            permits,
394            size,
395            io_metrics: OptIOMetrics(io_metrics),
396        })
397    }
398
399    pub fn set_io_metrics(&mut self, io_metrics: Option<Arc<IOMetrics>>) -> &mut Self {
400        self.io_metrics = OptIOMetrics(io_metrics);
401        self
402    }
403
404    pub fn io_metrics(&self) -> &OptIOMetrics {
405        &self.io_metrics
406    }
407}
408
409impl ByteSource for FileByteSource {
410    async fn get_size(&self) -> PolarsResult<usize> {
411        usize::try_from(self.size)
412            .map_err(|_| polars_err!(ComputeError: "file size {} does not fit in usize", self.size))
413    }
414
415    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
416        let size = self.get_size().await?;
417        let bytes = self.get_range(size.saturating_sub(n)..size).await?;
418        Ok((bytes, size))
419    }
420
421    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
422        assert!(range.end as u64 <= self.size);
423
424        let file = self.file.clone();
425        let offset = range.start as u64;
426        let len = range.len();
427        let size = self.size;
428        let o_direct = self.o_direct_align;
429
430        let permit = self.permits.clone().acquire_owned().await.unwrap();
431
432        self.io_metrics()
433            .record_io_read(len as u64, async move {
434                ASYNC
435                    .spawn_blocking(move || {
436                        let _permit = permit;
437                        read_at(&file, o_direct, offset, len, size)
438                    })
439                    .await
440            })
441            .await
442            .expect("blocking task panicked")
443    }
444
445    async fn get_ranges(
446        &self,
447        ranges: &mut [Range<usize>],
448    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
449        if let [range] = ranges {
450            let mut out = PlHashMap::with_capacity(1);
451            out.insert(range.start, self.get_range(range.clone()).await?);
452            return Ok(out);
453        }
454
455        ranges.sort_unstable_by_key(|r| r.start);
456
457        let mut spans: Vec<Range<usize>> = Vec::with_capacity(ranges.len());
458
459        // Threshold for coalescing. Note, individual ranges may exceed MAX_SPAN.
460        const MAX_SPAN: usize = 8 << 20;
461        // Tolerate small gaps. We match typical page size for now; this could be tuned.
462        const MAX_GAP: usize = 4096;
463
464        for r in ranges.iter() {
465            match spans.last_mut() {
466                Some(last)
467                    if r.start.saturating_sub(last.end) <= MAX_GAP
468                        && r.end.saturating_sub(last.start) <= MAX_SPAN =>
469                {
470                    last.end = last.end.max(r.end)
471                },
472                _ => spans.push(r.clone()),
473            }
474        }
475
476        let mut fetched: Vec<(Range<usize>, Buffer<u8>)> = futures::stream::iter(spans)
477            .map(|span| async move {
478                let buf = self.get_range(span.clone()).await?;
479                PolarsResult::Ok((span, buf))
480            })
481            .buffer_unordered(self.concurrency)
482            .try_collect()
483            .await?;
484
485        // Slice out of the containing span into the original ranges.
486        fetched.sort_unstable_by_key(|(s, _)| s.start);
487
488        if let Some(w) = ranges.windows(2).find(|w| w[0].start == w[1].start) {
489            polars_bail!(
490                ComputeError:
491                "duplicate range start {} in read request ({:?} and {:?})",
492                w[0].start, w[0], w[1]
493            );
494        }
495
496        let mut out = PlHashMap::with_capacity(ranges.len());
497        for r in ranges.iter() {
498            // No span available.
499            if r.is_empty() {
500                out.insert(r.start, Buffer::new());
501                continue;
502            }
503
504            // Spans are sorted by start and each range lies within one, so the
505            // containing span is the last one starting at or before `r.start`.
506            let idx = fetched.partition_point(|(s, _)| s.start <= r.start);
507            let (span, buf) = &fetched[idx - 1];
508            debug_assert!(span.start <= r.start && r.end <= span.end);
509
510            let off = r.start - span.start;
511            out.insert(r.start, buf.clone().sliced(off..off + r.len()));
512        }
513
514        debug_assert_eq!(out.len(), ranges.len());
515        Ok(out)
516    }
517}
518
519#[cfg(feature = "cloud")]
520pub struct ObjectStoreByteSource {
521    store: PolarsObjectStore,
522    path: ObjectStorePath,
523    config: FetchConfig,
524}
525
526#[cfg(feature = "cloud")]
527impl ObjectStoreByteSource {
528    async fn try_new_from_path(
529        path: PlRefPath,
530        cloud_options: Option<&CloudOptions>,
531        io_metrics: Option<Arc<IOMetrics>>,
532        config: FetchConfig,
533    ) -> PolarsResult<Self> {
534        let (CloudLocation { prefix, .. }, mut store) =
535            build_object_store(path, cloud_options, false).await?;
536        let path = object_path_from_str(&prefix)?;
537
538        store.set_io_metrics(io_metrics);
539
540        Ok(Self {
541            store,
542            path,
543            config,
544        })
545    }
546
547    #[allow(unused)]
548    fn chunk_size(&self) -> usize {
549        self.config.chunk_size
550    }
551
552    fn concurrency_strategy(&self) -> ConcurrencyStrategy {
553        self.config.strategy
554    }
555}
556
557#[cfg(feature = "cloud")]
558impl ByteSource for ObjectStoreByteSource {
559    async fn get_size(&self) -> PolarsResult<usize> {
560        Ok(self
561            .store
562            .head(&self.path, self.concurrency_strategy())
563            .await?
564            .size as usize)
565    }
566
567    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
568        self.store.get_suffix(&self.path, n, self.config).await
569    }
570
571    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
572        self.store.get_range(&self.path, range, self.config).await
573    }
574
575    async fn get_ranges(
576        &self,
577        ranges: &mut [Range<usize>],
578    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
579        self.store
580            .get_ranges_sort(&self.path, ranges, self.config)
581            .await
582    }
583}
584
585/// Dynamic dispatch to async functions.
586pub enum DynByteSource {
587    Buffer(BufferByteSource),
588    File(FileByteSource),
589    #[cfg(feature = "cloud")]
590    Cloud(ObjectStoreByteSource),
591}
592
593impl DynByteSource {
594    pub fn variant_name(&self) -> &str {
595        match self {
596            Self::Buffer(_) => "Buffer",
597            Self::File(_) => "File",
598            #[cfg(feature = "cloud")]
599            Self::Cloud(_) => "Cloud",
600        }
601    }
602
603    pub fn is_cloud(&self) -> bool {
604        match self {
605            Self::Buffer(_) => false,
606            Self::File(_) => false,
607            #[cfg(feature = "cloud")]
608            Self::Cloud(_) => true,
609        }
610    }
611
612    pub fn chunk_size(&self) -> Option<usize> {
613        match self {
614            Self::Buffer(_) => None,
615            Self::File(_) => None,
616            #[cfg(feature = "cloud")]
617            Self::Cloud(source) => Some(source.config.chunk_size),
618        }
619    }
620
621    pub fn concurrency_strategy(&self) -> Option<ConcurrencyStrategy> {
622        match self {
623            Self::Buffer(_) => None,
624            // Concurrency is handled directly inside FileByteSource.
625            Self::File(_) => None,
626            #[cfg(feature = "cloud")]
627            Self::Cloud(source) => Some(source.concurrency_strategy()),
628        }
629    }
630}
631
632impl Default for DynByteSource {
633    fn default() -> Self {
634        Self::Buffer(BufferByteSource(Buffer::new()))
635    }
636}
637
638impl ByteSource for DynByteSource {
639    async fn get_size(&self) -> PolarsResult<usize> {
640        match self {
641            Self::Buffer(v) => v.get_size().await,
642            Self::File(v) => v.get_size().await,
643            #[cfg(feature = "cloud")]
644            Self::Cloud(v) => v.get_size().await,
645        }
646    }
647
648    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
649        match self {
650            Self::Buffer(v) => v.get_suffix(n).await,
651            Self::File(v) => v.get_suffix(n).await,
652            #[cfg(feature = "cloud")]
653            Self::Cloud(v) => v.get_suffix(n).await,
654        }
655    }
656
657    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
658        match self {
659            Self::Buffer(v) => v.get_range(range).await,
660            Self::File(v) => v.get_range(range).await,
661            #[cfg(feature = "cloud")]
662            Self::Cloud(v) => v.get_range(range).await,
663        }
664    }
665
666    async fn get_ranges(
667        &self,
668        ranges: &mut [Range<usize>],
669    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
670        match self {
671            Self::Buffer(v) => v.get_ranges(ranges).await,
672            Self::File(v) => v.get_ranges(ranges).await,
673            #[cfg(feature = "cloud")]
674            Self::Cloud(v) => v.get_ranges(ranges).await,
675        }
676    }
677}
678
679impl From<BufferByteSource> for DynByteSource {
680    fn from(value: BufferByteSource) -> Self {
681        Self::Buffer(value)
682    }
683}
684
685impl From<FileByteSource> for DynByteSource {
686    fn from(value: FileByteSource) -> Self {
687        Self::File(value)
688    }
689}
690
691#[cfg(feature = "cloud")]
692impl From<ObjectStoreByteSource> for DynByteSource {
693    fn from(value: ObjectStoreByteSource) -> Self {
694        Self::Cloud(value)
695    }
696}
697
698impl From<Buffer<u8>> for DynByteSource {
699    fn from(value: Buffer<u8>) -> Self {
700        Self::Buffer(BufferByteSource(value))
701    }
702}
703
704#[derive(Clone, Debug)]
705pub enum DynByteSourceBuilder {
706    Mmap,
707    /// Use std::fs::File positional read (pread).
708    FilePread(FileReadContext),
709    /// Supports both cloud and local files, requires cloud feature.
710    ObjectStore(FetchConfig),
711}
712
713impl DynByteSourceBuilder {
714    pub async fn try_build_from_path(
715        &self,
716        path: PlRefPath,
717        cloud_options: Option<&CloudOptions>,
718        io_metrics: Option<Arc<IOMetrics>>,
719    ) -> PolarsResult<DynByteSource> {
720        Ok(match self {
721            Self::Mmap => {
722                BufferByteSource::try_new_mmap_from_path(path.as_std_path(), cloud_options)
723                    .await?
724                    .into()
725            },
726            Self::FilePread(read_context) => {
727                FileByteSource::try_new_from_path(path, read_context.clone(), io_metrics)
728                    .await?
729                    .into()
730            },
731            Self::ObjectStore(fetch_config) => feature_gated!("cloud", {
732                ObjectStoreByteSource::try_new_from_path(
733                    path,
734                    cloud_options,
735                    io_metrics,
736                    *fetch_config,
737                )
738                .await?
739                .into()
740            }),
741        })
742    }
743
744    pub fn chunk_size(&self) -> Option<usize> {
745        match self {
746            Self::Mmap => None,
747            Self::FilePread(_) => None,
748            Self::ObjectStore(fetch_config) => Some(fetch_config.chunk_size),
749        }
750    }
751
752    pub fn concurrency_strategy(&self) -> Option<&ConcurrencyStrategy> {
753        match self {
754            Self::Mmap => None,
755            Self::FilePread(_) => None,
756            Self::ObjectStore(fetch_config) => Some(&fetch_config.strategy),
757        }
758    }
759}
760
761#[cfg(test)]
762mod tests {
763    use super::*;
764
765    const MIB: usize = 1 << 20;
766
767    /// A file source over `len` bytes of `(i % 251)`, plus those bytes.
768    fn source(dir: &std::path::Path, len: usize) -> (FileByteSource, Vec<u8>) {
769        let path = dir.join("f.bin");
770        let contents: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
771        std::fs::write(&path, &contents).unwrap();
772
773        let read_context = FileReadContext {
774            enable_o_direct: false,
775            concurrency: 4,
776            permits: Arc::new(Semaphore::new(4)),
777            advice: FileAdvice::Normal,
778        };
779        let source = FileByteSource::try_new_from_std(
780            std::fs::File::open(&path).unwrap(),
781            read_context,
782            None,
783        )
784        .unwrap();
785        (source, contents)
786    }
787
788    fn runtime() -> tokio::runtime::Runtime {
789        tokio::runtime::Builder::new_multi_thread()
790            .worker_threads(2)
791            .enable_all()
792            .build()
793            .unwrap()
794    }
795
796    #[test]
797    fn get_ranges_resolves_every_range() {
798        let dir = tempfile::tempdir().unwrap();
799        let (source, contents) = source(dir.path(), 17 * MIB);
800        let rt = runtime();
801
802        let cases: [(&str, &[Range<usize>]); 6] = [
803            ("disjoint", &[0..100, 500..600, 900..1000]),
804            ("unsorted", &[900..1000, 0..100, 500..600]),
805            ("partial overlap", &[0..100, 50..150]),
806            ("contained", &[0..200, 50..100]),
807            // A zero-length column chunk is declarable in file metadata.
808            ("empty range at a span end", &[0..100, 100..100]),
809            // Spans become [0..4096, 5M..14M, 13M..16M]: the last two overlap,
810            // and both contain 13.5M, but only the third covers 16M.
811            (
812                "overlapping spans",
813                &[
814                    0..4096,
815                    5_000_000..14_000_000,
816                    13_000_000..13_500_000,
817                    13_500_000..16_000_000,
818                ],
819            ),
820        ];
821
822        for (name, ranges) in cases {
823            let mut ranges = ranges.to_vec();
824            let expect = ranges.clone();
825            let out = rt.block_on(source.get_ranges(&mut ranges)).unwrap();
826
827            assert_eq!(out.len(), expect.len(), "{name}");
828            for r in &expect {
829                let got = out
830                    .get(&r.start)
831                    .unwrap_or_else(|| panic!("{name}: {r:?} missing"));
832                assert_eq!(got.as_ref(), &contents[r.clone()], "{name}: {r:?}");
833            }
834        }
835    }
836}
polars_io/utils/byte_source.rs

polars_io/utils/
byte_source.rs
