Skip to main content

vortex_io/std_file/
read_at.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fs::File;
5use std::io;
6#[cfg(all(not(unix), not(windows)))]
7use std::io::Read;
8#[cfg(all(not(unix), not(windows)))]
9use std::io::Seek;
10#[cfg(unix)]
11use std::os::unix::fs::FileExt;
12#[cfg(windows)]
13use std::os::windows::fs::FileExt;
14use std::path::Path;
15use std::sync::Arc;
16
17use futures::FutureExt;
18use futures::future::BoxFuture;
19use vortex_array::buffer::BufferHandle;
20use vortex_array::memory::BufferAllocatorRef;
21use vortex_buffer::Alignment;
22use vortex_error::VortexResult;
23
24use crate::CoalesceConfig;
25use crate::VortexReadAt;
26use crate::runtime::Handle;
27
28/// Read exactly `buffer.len()` bytes from `file` starting at `offset`.
29/// This is a platform-specific helper that uses the most efficient method available.
30#[cfg(not(target_arch = "wasm32"))]
31pub fn read_exact_at(file: &File, buffer: &mut [u8], offset: u64) -> io::Result<()> {
32    #[cfg(unix)]
33    {
34        file.read_exact_at(buffer, offset)
35    }
36    #[cfg(windows)]
37    {
38        let mut bytes_read = 0;
39        while bytes_read < buffer.len() {
40            let read = file.seek_read(&mut buffer[bytes_read..], offset + bytes_read as u64)?;
41            if read == 0 {
42                return Err(io::Error::new(
43                    io::ErrorKind::UnexpectedEof,
44                    "failed to fill whole buffer",
45                ));
46            }
47            bytes_read += read;
48        }
49        Ok(())
50    }
51    #[cfg(all(not(unix), not(windows)))]
52    {
53        use std::io::SeekFrom;
54        let mut file_ref = file;
55        file_ref.seek(SeekFrom::Start(offset))?;
56        file_ref.read_exact(buffer)
57    }
58}
59
60/// Default number of concurrent requests to allow for local file I/O.
61pub const DEFAULT_CONCURRENCY: usize = 32;
62
63/// An adapter type wrapping a [`File`] to implement [`VortexReadAt`].
64pub struct FileReadAt {
65    uri: Arc<str>,
66    file: Arc<File>,
67    handle: Handle,
68    allocator: BufferAllocatorRef,
69}
70
71impl FileReadAt {
72    /// Open a file for reading.
73    pub fn open(path: impl AsRef<Path>, handle: Handle) -> VortexResult<Self> {
74        Self::open_with_allocator(path, handle, BufferAllocatorRef::statically_allocated())
75    }
76
77    /// Open a file for reading using a custom writable buffer allocator.
78    pub fn open_with_allocator(
79        path: impl AsRef<Path>,
80        handle: Handle,
81        allocator: BufferAllocatorRef,
82    ) -> VortexResult<Self> {
83        let path = path.as_ref();
84        let uri = path.to_string_lossy().to_string().into();
85        let file = Arc::new(File::open(path)?);
86        Ok(Self {
87            uri,
88            file,
89            handle,
90            allocator,
91        })
92    }
93}
94
95impl VortexReadAt for FileReadAt {
96    fn uri(&self) -> Option<&Arc<str>> {
97        Some(&self.uri)
98    }
99
100    fn coalesce_config(&self) -> Option<CoalesceConfig> {
101        Some(CoalesceConfig::file())
102    }
103
104    fn concurrency(&self) -> usize {
105        DEFAULT_CONCURRENCY
106    }
107
108    fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
109        let file = Arc::clone(&self.file);
110        async move {
111            let metadata = file.metadata()?;
112            Ok(metadata.len())
113        }
114        .boxed()
115    }
116
117    fn read_at(
118        &self,
119        offset: u64,
120        length: usize,
121        alignment: Alignment,
122    ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
123        let file = Arc::clone(&self.file);
124        let handle = self.handle.clone();
125        let allocator = self.allocator.clone();
126        async move {
127            handle
128                .spawn_blocking(move || {
129                    let mut buffer = allocator.with_capacity_aligned::<u8>(length, alignment);
130                    // SAFETY: read_exact_at initializes every byte before the buffer is frozen.
131                    unsafe { buffer.set_len(length) };
132                    read_exact_at(&file, buffer.as_mut_slice(), offset)?;
133                    Ok(BufferHandle::new_host(buffer.freeze()))
134                })
135                .await
136        }
137        .boxed()
138    }
139}