polars_io/utils/

byte_source.rs

use std::ops::Range;
#[cfg(target_os = "linux")]
use std::os::fd::AsRawFd;
use std::path::Path;
use std::sync::{Arc, LazyLock};

use dio_align::DioAlign;
use futures::{StreamExt, TryStreamExt};
use polars_buffer::Buffer;
use polars_config::FileAdvice;
use polars_core::prelude::PlHashMap;
use polars_core::runtime::ASYNC;
use polars_error::{PolarsResult, feature_gated, polars_bail, polars_err};
use polars_utils::aliases::InitHashMaps;
use polars_utils::io::_limit_path_len_io_err;
use polars_utils::mmap::MMapSemaphore;
use polars_utils::pl_path::PlRefPath;
use tokio::sync::Semaphore;

use crate::cloud::concurrency_config::{ConcurrencyStrategy, FetchConfig};
use crate::cloud::options::CloudOptions;
#[cfg(feature = "cloud")]
use crate::cloud::{
    CloudLocation, ObjectStorePath, PolarsObjectStore, build_object_store, object_path_from_str,
};
use crate::metrics::{IOMetrics, OptIOMetrics};

pub mod dio_align;
#[cfg(target_os = "linux")]
mod direct_io;
#[allow(async_fn_in_trait)]
pub trait ByteSource: Send + Sync {
    async fn get_size(&self) -> PolarsResult<usize>;
    /// Fetch the last `n` bytes and the total size of the source, in a single request. Returns
    /// fewer than `n` bytes if the source is smaller than `n`.
    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)>;
    /// Fetch the first `n` bytes and the total size of the source, in a single request where
    /// supported. Returns fewer than `n` bytes if the source is smaller than `n`.
    async fn get_prefix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)>;
    /// # Panics
    /// Panics if `range` is not in bounds.
    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>>;
    /// Note: This will mutably sort ranges for coalescing.
    async fn get_ranges(
        &self,
        ranges: &mut [Range<usize>],
    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>>;
}

/// Byte source backed by a `Buffer`, which can potentially be memory-mapped.
pub struct BufferByteSource(pub Buffer<u8>);

impl BufferByteSource {
    async fn try_new_mmap_from_path(
        path: &Path,
        _cloud_options: Option<&CloudOptions>,
    ) -> PolarsResult<Self> {
        let file = Arc::new(
            tokio::fs::File::open(path)
                .await
                .map_err(|err| _limit_path_len_io_err(path, err))?
                .into_std()
                .await,
        );

        Ok(Self(Buffer::from_owner(MMapSemaphore::new_from_file(
            &file,
        )?)))
    }
}

impl ByteSource for BufferByteSource {
    async fn get_size(&self) -> PolarsResult<usize> {
        Ok(self.0.as_ref().len())
    }

    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        let len = self.0.as_ref().len();
        Ok((self.0.clone().sliced(len.saturating_sub(n)..len), len))
    }

    async fn get_prefix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        let len = self.0.as_ref().len();
        Ok((self.0.clone().sliced(0..n.min(len)), len))
    }

    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
        let out = self.0.clone().sliced(range);
        Ok(out)
    }

    async fn get_ranges(
        &self,
        ranges: &mut [Range<usize>],
    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
        Ok(ranges
            .iter()
            .map(|x| (x.start, self.0.clone().sliced(x.clone())))
            .collect())
    }
}

/// Byte source backed by a `File`.
pub struct FileByteSource {
    file: Arc<std::fs::File>,
    // Alignment for O_DIRECT, if supported.
    o_direct_align: Option<DioAlign>,
    // Manage concurrency.
    concurrency: usize,
    permits: Arc<Semaphore>,
    // File size.
    size: u64,
    io_metrics: OptIOMetrics,
}

/// Each permit pins a tokio blocking thread for the duration of a `pread`.
pub fn global_read_permits() -> Arc<Semaphore> {
    static PERMITS: LazyLock<Arc<Semaphore>> = LazyLock::new(|| {
        Arc::new(Semaphore::new(
            polars_config::config().file_read_concurrency().max(1) as usize,
        ))
    });

    PERMITS.clone()
}

#[derive(Clone)]
pub struct FileReadContext {
    pub enable_o_direct: bool,
    pub concurrency: usize,
    pub permits: Arc<Semaphore>,
    /// Ignored when direct I/O is active: there is no page cache to advise.
    pub advice: FileAdvice,
}

impl FileReadContext {
    pub fn from_config(name: &str) -> Self {
        let cfg = polars_config::config();
        let enable_o_direct = cfg.direct_io();
        let concurrency = cfg.file_read_concurrency().max(1) as usize;

        // TODO: Posix_fadv should follow the access-pattern, which varies
        // by file type and projection.
        let fadv = cfg.file_posix_fadv();

        if cfg.verbose() {
            eprintln!(
                "[{name}]: file read_context as configured: \
                    read_concurrency: {concurrency}, \
                    posix_fadv: {fadv}, \
                    o_direct: {enable_o_direct}"
            );
        }

        Self {
            enable_o_direct,
            concurrency,
            permits: global_read_permits(),
            advice: fadv,
        }
    }
}

impl std::fmt::Debug for FileReadContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FileReadContext")
            .field("available_permits", &self.permits.available_permits())
            .field("enable_odirect", &self.enable_o_direct)
            .field("concurrency", &self.concurrency)
            .field("advice", &self.advice)
            .finish()
    }
}

#[cfg(target_os = "linux")]
fn _fadvise(file: &std::fs::File, advice: FileAdvice) {
    let advice = match advice {
        // The OS default is already POSIX_FADV_NORMAL.
        FileAdvice::Normal => return,
        FileAdvice::Sequential => libc::POSIX_FADV_SEQUENTIAL,
        FileAdvice::Random => libc::POSIX_FADV_RANDOM,
        FileAdvice::WillNeed => libc::POSIX_FADV_WILLNEED,
    };

    unsafe {
        libc::posix_fadvise(file.as_raw_fd(), 0, 0, advice);
    }
}

#[cfg(all(unix, not(feature = "nightly")))]
fn pread_exact(file: &std::fs::File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
    use std::os::unix::fs::FileExt;
    file.read_exact_at(buf, offset)
}

#[cfg(all(windows, not(feature = "nightly")))]
fn pread_exact(file: &std::fs::File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
    use std::os::windows::fs::FileExt;
    let mut filled = 0;
    while filled < buf.len() {
        match file.seek_read(&mut buf[filled..], offset + filled as u64)? {
            0 => return Err(std::io::ErrorKind::UnexpectedEof.into()),
            n => filled += n,
        }
    }
    Ok(())
}

/// `pread_exact` into uninitialized memory, so that `read_buffered` does not
/// have to zero the buffer the read is about to overwrite anyway.
#[cfg(all(feature = "nightly", unix))]
fn pread_exact_uninit(
    file: &std::fs::File,
    buf: std::io::BorrowedCursor<'_, u8>,
    offset: u64,
) -> std::io::Result<()> {
    use std::os::unix::fs::FileExt;
    file.read_buf_exact_at(buf, offset)
}

#[cfg(all(feature = "nightly", windows))]
fn pread_exact_uninit(
    file: &std::fs::File,
    mut buf: std::io::BorrowedCursor<'_, u8>,
    mut offset: u64,
) -> std::io::Result<()> {
    use std::os::windows::fs::FileExt;
    // `seek_read_buf` is the short-read variant; loop it to fill the cursor.
    while buf.capacity() > 0 {
        let written = buf.written();
        file.seek_read_buf(buf.reborrow(), offset)?;
        match buf.written() - written {
            0 => return Err(std::io::ErrorKind::UnexpectedEof.into()),
            n => offset += n as u64,
        }
    }
    Ok(())
}

/// Read `len` bytes at `offset` through the page cache.
///
/// Under `nightly` the destination is left uninitialized and filled by the read
/// itself; on stable it must be zeroed first, which memsets the whole range
/// before the kernel overwrites it.
#[cfg(feature = "nightly")]
fn read_buffered(file: &std::fs::File, offset: u64, len: usize) -> PolarsResult<Buffer<u8>> {
    let mut v: Vec<u8> = Vec::with_capacity(len);

    let filled = {
        let mut buf = std::io::BorrowedBuf::from(&mut v.spare_capacity_mut()[..len]);
        pread_exact_uninit(file, buf.unfilled(), offset)?;
        buf.len()
    };

    debug_assert_eq!(filled, len);
    // Safety: `pread_exact_uninit` filled the cursor, so the first `filled`
    // bytes of the allocation are initialized.
    unsafe { v.set_len(filled) };

    Ok(Buffer::from(v))
}

#[cfg(not(feature = "nightly"))]
fn read_buffered(file: &std::fs::File, offset: u64, len: usize) -> PolarsResult<Buffer<u8>> {
    let mut buf = vec![0u8; len];
    pread_exact(file, &mut buf, offset)?;
    Ok(Buffer::from(buf))
}

/// Read `len` bytes at `offset`, through direct I/O when `align` says the file
/// supports it and through the page cache otherwise.
#[cfg(target_os = "linux")]
fn read_at(
    file: &std::fs::File,
    align: Option<DioAlign>,
    offset: u64,
    len: usize,
    size: u64,
) -> PolarsResult<Buffer<u8>> {
    match align {
        Some(align) => direct_io::read_aligned(file, align, offset, len, size),
        None => read_buffered(file, offset, len),
    }
}

#[cfg(not(target_os = "linux"))]
fn read_at(
    file: &std::fs::File,
    _align: Option<DioAlign>,
    offset: u64,
    len: usize,
    _size: u64,
) -> PolarsResult<Buffer<u8>> {
    read_buffered(file, offset, len)
}

impl FileByteSource {
    async fn try_new_from_path(
        path: PlRefPath,
        read_context: FileReadContext,
        io_metrics: Option<Arc<IOMetrics>>,
    ) -> PolarsResult<Self> {
        // The open path is `open`, `fcntl`, `statx` and `fstat` - all blocking.
        ASYNC
            .spawn_blocking(move || Self::open_blocking(path, read_context, io_metrics))
            .await
            .expect("blocking task panicked")
    }

    fn open_blocking(
        path: PlRefPath,
        read_context: FileReadContext,
        io_metrics: Option<Arc<IOMetrics>>,
    ) -> PolarsResult<Self> {
        let path = path.as_std_path();
        let enable_o_direct = read_context.enable_o_direct;

        let file = {
            #[cfg(target_os = "linux")]
            let f = if enable_o_direct {
                direct_io::open_o_direct(path).or_else(|e| match e.raw_os_error() {
                    // The filesystem does not support O_DIRECT.
                    Some(libc::EINVAL) => std::fs::File::open(path),
                    _ => Err(e),
                })
            } else {
                std::fs::File::open(path)
            };
            #[cfg(not(target_os = "linux"))]
            let f = std::fs::File::open(path);

            f.map_err(|e| _limit_path_len_io_err(path, e))?
        };

        let file = Arc::new(file);

        #[cfg(target_os = "linux")]
        let o_direct_align =
            direct_io::probe_or_disable(&file).map_err(|e| _limit_path_len_io_err(path, e))?;

        #[cfg(not(target_os = "linux"))]
        let o_direct_align: Option<DioAlign> = None;

        #[cfg(target_os = "linux")]
        if o_direct_align.is_none() {
            // Inert under O_DIRECT: there is no page cache to advise.
            _fadvise(&file, read_context.advice);
        }

        let concurrency = read_context.concurrency;
        let permits = read_context.permits;

        let size = file
            .metadata()
            .map_err(|e| _limit_path_len_io_err(path, e))?
            .len();

        if polars_config::config().verbose() {
            let name = if polars_config::config().verbose_sensitive() {
                path.display().to_string()
            } else {
                "<file>".to_string()
            };
            match (enable_o_direct, o_direct_align) {
                (true, Some(a)) => eprintln!(
                    "[FileByteSource]: {name}: direct IO active, \
                        alignment offset: {}, memory: {}",
                    a.offset, a.memory
                ),
                (true, None) => eprintln!(
                    "[FileByteSource]: {name}: direct IO requested but not active, \
                        using buffered reads"
                ),
                (false, _) => {},
            }
        }

        Ok(FileByteSource {
            file,
            o_direct_align,
            concurrency,
            permits,
            size,
            io_metrics: OptIOMetrics(io_metrics),
        })
    }

    pub fn try_new_from_std(
        file: std::fs::File,
        read_context: FileReadContext,
        io_metrics: Option<Arc<IOMetrics>>,
    ) -> PolarsResult<Self> {
        let size = file.metadata()?.len();

        #[cfg(target_os = "linux")]
        let o_direct_align = direct_io::probe_or_disable(&file)?;

        #[cfg(not(target_os = "linux"))]
        let o_direct_align: Option<DioAlign> = None;

        #[cfg(target_os = "linux")]
        if o_direct_align.is_none() {
            _fadvise(&file, read_context.advice);
        }

        let concurrency = read_context.concurrency;
        let permits = read_context.permits;

        // For verbose logging only. We cannot enable since the file_handle was handed to us.
        let enable_o_direct = read_context.enable_o_direct;

        if polars_config::config().verbose() {
            match (enable_o_direct, o_direct_align) {
                (true, Some(a)) => eprintln!(
                    "[FileByteSource]: direct IO active, alignment offset: {}, memory: {}",
                    a.offset, a.memory
                ),
                (true, None) => eprintln!(
                    "[FileByteSource]: direct IO requested but not active, using buffered reads",
                ),
                (false, _) => {},
            }
        }

        Ok(Self {
            file: Arc::new(file),
            o_direct_align,
            concurrency,
            permits,
            size,
            io_metrics: OptIOMetrics(io_metrics),
        })
    }

    pub fn set_io_metrics(&mut self, io_metrics: Option<Arc<IOMetrics>>) -> &mut Self {
        self.io_metrics = OptIOMetrics(io_metrics);
        self
    }

    pub fn io_metrics(&self) -> &OptIOMetrics {
        &self.io_metrics
    }

    /// Whether every page of `ranges` is in the page cache. `false` when the OS cannot tell
    /// (non-Linux, kernels before 6.5) and under direct I/O.
    pub fn is_cached(&self, ranges: &[Range<usize>]) -> bool {
        #[cfg(all(
            target_os = "linux",
            any(target_arch = "x86_64", target_arch = "aarch64")
        ))]
        if self.o_direct_align.is_none() {
            return cachestat::all_cached(&self.file, ranges);
        }

        let _ = ranges;
        false
    }

    /// Reads `range` on the calling thread.
    pub fn read_blocking(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
        assert!(range.end as u64 <= self.size);

        self.io_metrics()
            .record_io_read_blocking(range.len() as u64, || {
                read_at(
                    &self.file,
                    self.o_direct_align,
                    range.start as u64,
                    range.len(),
                    self.size,
                )
            })
    }
}

/// `cachestat(2)`, Linux 6.5+.
#[cfg(all(
    target_os = "linux",
    any(target_arch = "x86_64", target_arch = "aarch64")
))]
mod cachestat {
    use std::ops::Range;
    use std::os::fd::AsRawFd;
    use std::sync::atomic::{AtomicBool, Ordering};

    use polars_utils::mem::PAGE_SIZE;

    #[repr(C)]
    struct CachestatRange {
        off: u64,
        len: u64,
    }

    #[repr(C)]
    #[derive(Default)]
    struct Cachestat {
        nr_cache: u64,
        nr_dirty: u64,
        nr_writeback: u64,
        nr_evicted: u64,
        nr_recently_evicted: u64,
    }

    /// Not in `libc` for these targets yet; the same number on both.
    const SYS_CACHESTAT: libc::c_long = 451;

    /// Set once the kernel has refused the call, so later checks skip the syscall.
    static UNSUPPORTED: AtomicBool = AtomicBool::new(false);

    /// Whether every page overlapping `ranges` is in the page cache.
    pub(super) fn all_cached(file: &std::fs::File, ranges: &[Range<usize>]) -> bool {
        if UNSUPPORTED.load(Ordering::Relaxed) {
            return false;
        }

        // One call per run of adjacent or overlapping ranges, e.g. the column chunks of a row
        // group. Gaps are not bridged: their pages need not be cached.
        let mut spans: Vec<Range<usize>> =
            ranges.iter().filter(|r| !r.is_empty()).cloned().collect();
        spans.sort_unstable_by_key(|r| r.start);
        spans.dedup_by(|next, cur| {
            let merge = next.start <= cur.end;
            if merge {
                cur.end = cur.end.max(next.end);
            }
            merge
        });

        let page = *PAGE_SIZE;
        for r in &spans {
            let mut range = CachestatRange {
                off: r.start as u64,
                len: r.len() as u64,
            };
            let mut cs = Cachestat::default();
            // Safety: both structs match the kernel's layout and outlive the call.
            let ret = unsafe {
                libc::syscall(
                    SYS_CACHESTAT,
                    file.as_raw_fd() as libc::c_long,
                    &mut range as *mut CachestatRange,
                    &mut cs as *mut Cachestat,
                    0 as libc::c_long,
                )
            };
            if ret != 0 {
                // ENOSYS before 6.5 and EPERM under some seccomp filters hold for every file.
                // Others, like EOPNOTSUPP on hugetlbfs, are about this file only.
                if let Some(libc::ENOSYS | libc::EPERM) =
                    std::io::Error::last_os_error().raw_os_error()
                {
                    UNSUPPORTED.store(true, Ordering::Relaxed);
                }
                return false;
            }
            // The kernel counts the pages that overlap the byte range.
            let pages = (r.end - 1) / page - r.start / page + 1;
            if cs.nr_cache < pages as u64 {
                return false;
            }
        }
        true
    }
}

impl ByteSource for FileByteSource {
    async fn get_size(&self) -> PolarsResult<usize> {
        usize::try_from(self.size)
            .map_err(|_| polars_err!(ComputeError: "file size {} does not fit in usize", self.size))
    }

    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        let size = self.get_size().await?;
        let bytes = self.get_range(size.saturating_sub(n)..size).await?;
        Ok((bytes, size))
    }

    async fn get_prefix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        let size = self.get_size().await?;
        let bytes = self.get_range(0..n.min(size)).await?;
        Ok((bytes, size))
    }

    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
        assert!(range.end as u64 <= self.size);

        let file = self.file.clone();
        let offset = range.start as u64;
        let len = range.len();
        let size = self.size;
        let o_direct = self.o_direct_align;

        let permit = self.permits.clone().acquire_owned().await.unwrap();

        self.io_metrics()
            .record_io_read(len as u64, async move {
                ASYNC
                    .spawn_blocking(move || {
                        let _permit = permit;
                        read_at(&file, o_direct, offset, len, size)
                    })
                    .await
            })
            .await
            .expect("blocking task panicked")
    }

    async fn get_ranges(
        &self,
        ranges: &mut [Range<usize>],
    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
        if let [range] = ranges {
            let mut out = PlHashMap::with_capacity(1);
            out.insert(range.start, self.get_range(range.clone()).await?);
            return Ok(out);
        }

        ranges.sort_unstable_by_key(|r| r.start);

        let mut spans: Vec<Range<usize>> = Vec::with_capacity(ranges.len());

        // Threshold for coalescing. Note, individual ranges may exceed MAX_SPAN.
        const MAX_SPAN: usize = 8 << 20;
        // Tolerate small gaps. We match typical page size for now; this could be tuned.
        const MAX_GAP: usize = 4096;

        for r in ranges.iter() {
            match spans.last_mut() {
                Some(last)
                    if r.start.saturating_sub(last.end) <= MAX_GAP
                        && r.end.saturating_sub(last.start) <= MAX_SPAN =>
                {
                    last.end = last.end.max(r.end)
                },
                _ => spans.push(r.clone()),
            }
        }

        let mut fetched: Vec<(Range<usize>, Buffer<u8>)> = futures::stream::iter(spans)
            .map(|span| async move {
                let buf = self.get_range(span.clone()).await?;
                PolarsResult::Ok((span, buf))
            })
            .buffer_unordered(self.concurrency)
            .try_collect()
            .await?;

        // Slice out of the containing span into the original ranges.
        fetched.sort_unstable_by_key(|(s, _)| s.start);

        if let Some(w) = ranges.windows(2).find(|w| w[0].start == w[1].start) {
            polars_bail!(
                ComputeError:
                "duplicate range start {} in read request ({:?} and {:?})",
                w[0].start, w[0], w[1]
            );
        }

        let mut out = PlHashMap::with_capacity(ranges.len());
        for r in ranges.iter() {
            // No span available.
            if r.is_empty() {
                out.insert(r.start, Buffer::new());
                continue;
            }

            // Spans are sorted by start and each range lies within one, so the
            // containing span is the last one starting at or before `r.start`.
            let idx = fetched.partition_point(|(s, _)| s.start <= r.start);
            let (span, buf) = &fetched[idx - 1];
            debug_assert!(span.start <= r.start && r.end <= span.end);

            let off = r.start - span.start;
            out.insert(r.start, buf.clone().sliced(off..off + r.len()));
        }

        debug_assert_eq!(out.len(), ranges.len());
        Ok(out)
    }
}

#[cfg(feature = "cloud")]
#[derive(Clone)]
pub struct ObjectStoreByteSource {
    store: PolarsObjectStore,
    path: ObjectStorePath,
    config: FetchConfig,
}

#[cfg(feature = "cloud")]
impl ObjectStoreByteSource {
    async fn try_new_from_path(
        path: PlRefPath,
        cloud_options: Option<&CloudOptions>,
        io_metrics: Option<Arc<IOMetrics>>,
        config: FetchConfig,
    ) -> PolarsResult<Self> {
        let (CloudLocation { prefix, .. }, mut store) =
            build_object_store(path, cloud_options, false).await?;
        let path = object_path_from_str(&prefix)?;

        store.set_io_metrics(io_metrics);

        Ok(Self {
            store,
            path,
            config,
        })
    }

    #[allow(unused)]
    fn chunk_size(&self) -> usize {
        self.config.chunk_size
    }

    fn concurrency_strategy(&self) -> ConcurrencyStrategy {
        self.config.strategy
    }
}

#[cfg(feature = "cloud")]
impl ByteSource for ObjectStoreByteSource {
    async fn get_size(&self) -> PolarsResult<usize> {
        Ok(self
            .store
            .head(&self.path, self.concurrency_strategy())
            .await?
            .size as usize)
    }

    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        self.store.get_suffix(&self.path, n, self.config).await
    }

    async fn get_prefix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        self.store.get_prefix(&self.path, n, self.config).await
    }

    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
        self.store.get_range(&self.path, range, self.config).await
    }

    async fn get_ranges(
        &self,
        ranges: &mut [Range<usize>],
    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
        self.store
            .get_ranges_sort(&self.path, ranges, self.config)
            .await
    }
}

/// Dynamic dispatch to async functions.
pub enum DynByteSource {
    Buffer(BufferByteSource),
    File(FileByteSource),
    #[cfg(feature = "cloud")]
    Cloud(ObjectStoreByteSource),
}

impl DynByteSource {
    pub fn variant_name(&self) -> &str {
        match self {
            Self::Buffer(_) => "Buffer",
            Self::File(_) => "File",
            #[cfg(feature = "cloud")]
            Self::Cloud(_) => "Cloud",
        }
    }

    /// A clone of this object store source with a different fetch configuration. The clone
    /// shares the underlying store, and this source keeps its configuration. Returns `None` for
    /// non-object store sources.
    #[cfg_attr(not(feature = "cloud"), allow(unused_variables))]
    pub fn with_fetch_config(&self, config: FetchConfig) -> Option<Self> {
        match self {
            Self::Buffer(_) | Self::File(_) => None,
            #[cfg(feature = "cloud")]
            Self::Cloud(v) => Some(Self::Cloud(ObjectStoreByteSource {
                config,
                ..v.clone()
            })),
        }
    }

    pub fn is_cloud(&self) -> bool {
        match self {
            Self::Buffer(_) => false,
            Self::File(_) => false,
            #[cfg(feature = "cloud")]
            Self::Cloud(_) => true,
        }
    }

    pub fn chunk_size(&self) -> Option<usize> {
        match self {
            Self::Buffer(_) => None,
            Self::File(_) => None,
            #[cfg(feature = "cloud")]
            Self::Cloud(source) => Some(source.config.chunk_size),
        }
    }

    pub fn concurrency_strategy(&self) -> Option<ConcurrencyStrategy> {
        match self {
            Self::Buffer(_) => None,
            // Concurrency is handled directly inside FileByteSource.
            Self::File(_) => None,
            #[cfg(feature = "cloud")]
            Self::Cloud(source) => Some(source.concurrency_strategy()),
        }
    }
}

impl Default for DynByteSource {
    fn default() -> Self {
        Self::Buffer(BufferByteSource(Buffer::new()))
    }
}

impl ByteSource for DynByteSource {
    async fn get_size(&self) -> PolarsResult<usize> {
        match self {
            Self::Buffer(v) => v.get_size().await,
            Self::File(v) => v.get_size().await,
            #[cfg(feature = "cloud")]
            Self::Cloud(v) => v.get_size().await,
        }
    }

    async fn get_suffix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        match self {
            Self::Buffer(v) => v.get_suffix(n).await,
            Self::File(v) => v.get_suffix(n).await,
            #[cfg(feature = "cloud")]
            Self::Cloud(v) => v.get_suffix(n).await,
        }
    }

    async fn get_prefix(&self, n: usize) -> PolarsResult<(Buffer<u8>, usize)> {
        match self {
            Self::Buffer(v) => v.get_prefix(n).await,
            Self::File(v) => v.get_prefix(n).await,
            #[cfg(feature = "cloud")]
            Self::Cloud(v) => v.get_prefix(n).await,
        }
    }

    async fn get_range(&self, range: Range<usize>) -> PolarsResult<Buffer<u8>> {
        match self {
            Self::Buffer(v) => v.get_range(range).await,
            Self::File(v) => v.get_range(range).await,
            #[cfg(feature = "cloud")]
            Self::Cloud(v) => v.get_range(range).await,
        }
    }

    async fn get_ranges(
        &self,
        ranges: &mut [Range<usize>],
    ) -> PolarsResult<PlHashMap<usize, Buffer<u8>>> {
        match self {
            Self::Buffer(v) => v.get_ranges(ranges).await,
            Self::File(v) => v.get_ranges(ranges).await,
            #[cfg(feature = "cloud")]
            Self::Cloud(v) => v.get_ranges(ranges).await,
        }
    }
}

impl From<BufferByteSource> for DynByteSource {
    fn from(value: BufferByteSource) -> Self {
        Self::Buffer(value)
    }
}

impl From<FileByteSource> for DynByteSource {
    fn from(value: FileByteSource) -> Self {
        Self::File(value)
    }
}

#[cfg(feature = "cloud")]
impl From<ObjectStoreByteSource> for DynByteSource {
    fn from(value: ObjectStoreByteSource) -> Self {
        Self::Cloud(value)
    }
}

impl From<Buffer<u8>> for DynByteSource {
    fn from(value: Buffer<u8>) -> Self {
        Self::Buffer(BufferByteSource(value))
    }
}

#[derive(Clone, Debug)]
pub enum DynByteSourceBuilder {
    Mmap,
    /// Use std::fs::File positional read (pread).
    FilePread(FileReadContext),
    /// Supports both cloud and local files, requires cloud feature.
    ObjectStore(FetchConfig),
}

impl DynByteSourceBuilder {
    pub async fn try_build_from_path(
        &self,
        path: PlRefPath,
        cloud_options: Option<&CloudOptions>,
        io_metrics: Option<Arc<IOMetrics>>,
    ) -> PolarsResult<DynByteSource> {
        Ok(match self {
            Self::Mmap => {
                BufferByteSource::try_new_mmap_from_path(path.as_std_path(), cloud_options)
                    .await?
                    .into()
            },
            Self::FilePread(read_context) => {
                FileByteSource::try_new_from_path(path, read_context.clone(), io_metrics)
                    .await?
                    .into()
            },
            Self::ObjectStore(fetch_config) => feature_gated!("cloud", {
                ObjectStoreByteSource::try_new_from_path(
                    path,
                    cloud_options,
                    io_metrics,
                    *fetch_config,
                )
                .await?
                .into()
            }),
        })
    }

    pub fn chunk_size(&self) -> Option<usize> {
        match self {
            Self::Mmap => None,
            Self::FilePread(_) => None,
            Self::ObjectStore(fetch_config) => Some(fetch_config.chunk_size),
        }
    }

    pub fn concurrency_strategy(&self) -> Option<&ConcurrencyStrategy> {
        match self {
            Self::Mmap => None,
            Self::FilePread(_) => None,
            Self::ObjectStore(fetch_config) => Some(&fetch_config.strategy),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const MIB: usize = 1 << 20;

    /// A file source over `len` bytes of `(i % 251)`, plus those bytes.
    fn source(dir: &std::path::Path, len: usize) -> (FileByteSource, Vec<u8>) {
        let path = dir.join("f.bin");
        let contents: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
        std::fs::write(&path, &contents).unwrap();

        let read_context = FileReadContext {
            enable_o_direct: false,
            concurrency: 4,
            permits: Arc::new(Semaphore::new(4)),
            advice: FileAdvice::Normal,
        };
        let source = FileByteSource::try_new_from_std(
            std::fs::File::open(&path).unwrap(),
            read_context,
            None,
        )
        .unwrap();
        (source, contents)
    }

    fn runtime() -> tokio::runtime::Runtime {
        tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .unwrap()
    }

    #[test]
    fn get_ranges_resolves_every_range() {
        let dir = tempfile::tempdir().unwrap();
        let (source, contents) = source(dir.path(), 17 * MIB);
        let rt = runtime();

        let cases: [(&str, &[Range<usize>]); 6] = [
            ("disjoint", &[0..100, 500..600, 900..1000]),
            ("unsorted", &[900..1000, 0..100, 500..600]),
            ("partial overlap", &[0..100, 50..150]),
            ("contained", &[0..200, 50..100]),
            // A zero-length column chunk is declarable in file metadata.
            ("empty range at a span end", &[0..100, 100..100]),
            // Spans become [0..4096, 5M..14M, 13M..16M]: the last two overlap,
            // and both contain 13.5M, but only the third covers 16M.
            (
                "overlapping spans",
                &[
                    0..4096,
                    5_000_000..14_000_000,
                    13_000_000..13_500_000,
                    13_500_000..16_000_000,
                ],
            ),
        ];

        for (name, ranges) in cases {
            let mut ranges = ranges.to_vec();
            let expect = ranges.clone();
            let out = rt.block_on(source.get_ranges(&mut ranges)).unwrap();

            assert_eq!(out.len(), expect.len(), "{name}");
            for r in &expect {
                let got = out
                    .get(&r.start)
                    .unwrap_or_else(|| panic!("{name}: {r:?} missing"));
                assert_eq!(got.as_ref(), &contents[r.clone()], "{name}: {r:?}");
            }
        }
    }

    #[test]
    fn read_blocking_reads_the_range() {
        let dir = tempfile::tempdir().unwrap();
        let (source, contents) = source(dir.path(), 3 * MIB);

        for r in [0..0, 0..1, 100..4096, MIB - 7..MIB + 9, 0..3 * MIB] {
            let got = source.read_blocking(r.clone()).unwrap();
            assert_eq!(got.as_ref(), &contents[r.clone()], "{r:?}");
        }
    }

    #[cfg(all(
        target_os = "linux",
        any(target_arch = "x86_64", target_arch = "aarch64")
    ))]
    #[test]
    fn is_cached_follows_the_page_cache() {
        let dir = tempfile::tempdir().unwrap();
        let (source, _) = source(dir.path(), 4 * MIB);
        let ranges = [0..100, MIB + 1..2 * MIB + 1, 4 * MIB - 1..4 * MIB, 5..5];

        // Just written, so cached; `false` means the kernel lacks cachestat(2).
        if !source.is_cached(&ranges) {
            return;
        }
        assert!(source.is_cached(&[]));

        // Drop the pages. tmpfs keeps them, so only check where they can go.
        source.file.sync_all().unwrap();
        unsafe { libc::posix_fadvise(source.file.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED) };
        let mut st: libc::statfs = unsafe { std::mem::zeroed() };
        let path = std::ffi::CString::new(dir.path().as_os_str().as_encoded_bytes()).unwrap();
        assert_eq!(unsafe { libc::statfs(path.as_ptr(), &mut st) }, 0);
        if st.f_type == libc::TMPFS_MAGIC {
            return;
        }
        assert!(!source.is_cached(&ranges));

        let read = MIB + 1..2 * MIB + 1;
        source.read_blocking(read.clone()).unwrap();
        assert!(source.is_cached(std::slice::from_ref(&read)));
        assert!(!source.is_cached(&ranges));
    }

    #[cfg(all(
        target_os = "linux",
        any(target_arch = "x86_64", target_arch = "aarch64")
    ))]
    #[test]
    fn is_cached_does_not_bridge_gaps() {
        let page = *polars_utils::mem::PAGE_SIZE;
        let dir = tempfile::tempdir().unwrap();
        let (source, _) = source(dir.path(), 3 * page);
        let (all, middle) = (0..3 * page, page..2 * page);
        // Just written, so cached; `false` means the kernel lacks cachestat(2).
        if !source.is_cached(std::slice::from_ref(&all)) {
            return;
        }

        // Drop only the middle page. tmpfs keeps it, so only check where it can go.
        source.file.sync_all().unwrap();
        unsafe {
            libc::posix_fadvise(
                source.file.as_raw_fd(),
                page as libc::off_t,
                page as libc::off_t,
                libc::POSIX_FADV_DONTNEED,
            )
        };
        if source.is_cached(std::slice::from_ref(&middle)) {
            return;
        }
        assert!(source.is_cached(&[2 * page..3 * page, 0..page]));
        assert!(!source.is_cached(&[0..page, page..2 * page]));
        assert!(!source.is_cached(std::slice::from_ref(&all)));
    }
}