From ee047209fec63f177d993ee9e36bfa966dfddcb9 Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Wed, 22 Jul 2026 13:55:45 +0100 Subject: [PATCH 1/8] feat[cuda]: add direct I/O for local file scans Keep footer and zone-map reads buffered while data-plane segments are read into pooled pinned buffers with Linux direct I/O. Expose the mode through both CUDA open options and a backwards-compatible FFI scan options entry point. Signed-off-by: Onur Satici --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-cuda/Cargo.toml | 3 + vortex-cuda/ffi/README.md | 4 + vortex-cuda/ffi/cinclude/vortex_cuda.h | 27 +++ vortex-cuda/ffi/src/lib.rs | 78 +++++++- vortex-cuda/src/file.rs | 27 ++- vortex-cuda/src/pinned.rs | 10 + vortex-cuda/src/pooled_read_at.rs | 251 ++++++++++++++++++++++++- 9 files changed, 389 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ab869e730b9..33f17258b95 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9959,6 +9959,7 @@ dependencies = [ "parking_lot", "prost 0.14.4", "rstest", + "rustix", "tokio", "tracing", "vortex", diff --git a/Cargo.toml b/Cargo.toml index 7581d1c9653..c25713c0222 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -230,6 +230,7 @@ roaring = "0.11.0" rstest = "0.26.1" rstest_reuse = "0.7.0" rustc-hash = "2.1.1" +rustix = { version = "1.1", features = ["fs"] } serde = "1.0.220" serde_json = "1.0.138" serde_test = "1.0.176" diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index be6faa77e7b..a0c57478ce1 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -44,6 +44,9 @@ vortex-cuda-macros = { workspace = true } vortex-error = { workspace = true, features = ["object_store"] } vortex-nvcomp = { path = "nvcomp" } +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true } + [dev-dependencies] criterion = { workspace = true } futures = { workspace = true, features = ["executor"] } diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 1b07e12f0df..3a6098c0342 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -23,3 +23,7 @@ Push host-resident arrays and close the sink using the standard `vx_array_sink_* Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer pool and CUDA state are reused as well. + +On Linux, use `vx_cuda_scan_path_arrow_device_stream_with_options` with +`vx_cuda_scan_options.direct_io = true` to bypass the operating system page cache for data-plane +reads. Footer and zone-map reads remain buffered on the host. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index a7c080e13cb..6f1e171c739 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -94,6 +94,19 @@ vx_array_sink *vx_cuda_array_sink_open_file_block_rows(const vx_session *session size_t block_rows, vx_error **error_out); +/** + * Options for scanning a CUDA-compatible Vortex file. + * + * Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. + */ +typedef struct vx_cuda_scan_options { + /** Bypass the operating system page cache for data-plane reads. + * Footer and zone-map reads remain buffered. Supported only on Linux. */ + bool direct_io; + /** Number of rows in each output ArrowDeviceArray. Zero uses layout-derived splitting. */ + size_t batch_rows; +} vx_cuda_scan_options; + /** * Scan a local CUDA-compatible Vortex file as an Arrow C Device stream. * @@ -125,6 +138,20 @@ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, struct ArrowDeviceArrayStream *out_stream, vx_error **error_out); +/** + * Scan a local CUDA-compatible Vortex file with explicit options. + * + * This has the same ownership and file compatibility requirements as + * `vx_cuda_scan_path_arrow_device_stream`. Pass NULL or a zero-initialized options struct to use + * buffered file I/O and layout-derived batch splitting. + */ +int vx_cuda_scan_path_arrow_device_stream_with_options( + const vx_session *session, + vx_view path, + const vx_cuda_scan_options *options, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index be2a9defc59..abe6baa3a15 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -49,6 +49,20 @@ use vortex_ffi::vx_view; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; +/// Options for scanning a CUDA-compatible Vortex file. +/// +/// Zero-initialize this struct to use buffered file I/O and layout-derived batch splitting. +#[repr(C)] +#[derive(Default)] +pub struct vx_cuda_scan_options { + /// Bypass the operating system page cache for data-plane reads. + /// + /// Footer and zone-map reads remain buffered. Direct I/O is supported only on Linux. + pub direct_io: bool, + /// Number of rows in each output batch. Zero uses layout-derived splitting. + pub batch_rows: usize, +} + /// Return a Vortex session with a [`CudaSession`] session variable. /// /// If `session` already has CUDA support, this returns a clone of it. Otherwise it @@ -142,7 +156,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( }) } -/// Scan a local Vortex file and export an Arrow C Device stream. +/// Scan a local Vortex file with buffered I/O and export an Arrow C Device stream. /// /// Footer and zone-map reads remain on the host. Data segments are staged through pinned host /// buffers and transferred directly to the GPU. @@ -171,7 +185,13 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( error_out: *mut *mut vx_error, ) -> c_int { unsafe { - vx_cuda_scan_path_arrow_device_stream_batch_rows(session, path, 0, out_stream, error_out) + vx_cuda_scan_path_arrow_device_stream_with_options( + session, + path, + ptr::null(), + out_stream, + error_out, + ) } } @@ -195,14 +215,59 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_batch_rows batch_rows: usize, out_stream: *mut ArrowDeviceArrayStream, error_out: *mut *mut vx_error, +) -> c_int { + let options = vx_cuda_scan_options { + batch_rows, + ..Default::default() + }; + unsafe { + vx_cuda_scan_path_arrow_device_stream_with_options( + session, + path, + &raw const options, + out_stream, + error_out, + ) + } +} + +/// Scan a local Vortex file with explicit options and export an Arrow C Device stream. +/// +/// This has the same ownership and file compatibility requirements as +/// [`vx_cuda_scan_path_arrow_device_stream`]. Pass a null `options` pointer or a zero-initialized +/// [`vx_cuda_scan_options`] to use buffered file I/O and layout-derived batch splitting. +/// +/// # Safety +/// +/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the +/// duration of this call and contain UTF-8. `options`, when non-null, must point to a valid +/// [`vx_cuda_scan_options`]. `out_stream` must be a valid writable pointer. If `error_out` is +/// non-null, it must be valid for writing one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream_with_options( + session: *const vx_session, + path: vx_view, + options: *const vx_cuda_scan_options, + out_stream: *mut ArrowDeviceArrayStream, + error_out: *mut *mut vx_error, ) -> c_int { try_or(error_out, VX_CUDA_ERR, || { vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output"); let path = unsafe { path.as_str() }?.to_owned(); let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; + let (direct_io, batch_rows) = if options.is_null() { + (false, 0) + } else { + let options = unsafe { &*options }; + (options.direct_io, options.batch_rows) + }; let array_stream = ffi_runtime().block_on(async { - let file = session.open_options().with_cuda().open_path(path).await?; + let mut options = session.open_options().with_cuda(); + if direct_io { + options = options.with_direct_io(); + } + let file = options.open_path(path).await?; let scan = file.scan()?; let scan = if batch_rows == 0 { scan @@ -321,6 +386,13 @@ mod tests { use super::*; + #[test] + fn scan_options_default_to_buffered_io() { + let options = vx_cuda_scan_options::default(); + assert!(!options.direct_io); + assert_eq!(options.batch_rows, 0); + } + fn test_session(session: VortexSession) -> *mut vx_session { Box::into_raw(Box::new(session)).cast::() } diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs index 9a41f51f9b9..76e033af856 100644 --- a/vortex-cuda/src/file.rs +++ b/vortex-cuda/src/file.rs @@ -31,7 +31,10 @@ pub trait CudaOpenOptionsExt { impl CudaOpenOptionsExt for VortexOpenOptions { fn with_cuda(self) -> CudaOpenOptions { - CudaOpenOptions { inner: self } + CudaOpenOptions { + inner: self, + direct_io: false, + } } } @@ -41,9 +44,19 @@ impl CudaOpenOptionsExt for VortexOpenOptions { /// configured before calling `with_cuda`. pub struct CudaOpenOptions { inner: VortexOpenOptions, + direct_io: bool, } impl CudaOpenOptions { + /// Read data-plane segments with direct I/O, bypassing the operating system page cache. + /// + /// This option is currently supported only on Linux. Footer and zone-map reads continue to + /// use buffered host I/O. + pub fn with_direct_io(mut self) -> Self { + self.direct_io = true; + self + } + /// Open a local Vortex file for CUDA execution. /// /// The footer and zone-map segments are read through the ordinary host path. All other file @@ -65,12 +78,12 @@ impl CudaOpenOptions { let pool = Arc::clone(cuda_session.pinned_buffer_pool()); drop(cuda_session); - let reader = Arc::new(PooledFileReadAt::open( - &path, - session.handle(), - pool, - stream, - )?); + let reader = if self.direct_io { + PooledFileReadAt::open_direct(&path, session.handle(), pool, stream)? + } else { + PooledFileReadAt::open(&path, session.handle(), pool, stream)? + }; + let reader = Arc::new(reader); let data_file = data_options .with_footer(footer.clone()) .open(reader) diff --git a/vortex-cuda/src/pinned.rs b/vortex-cuda/src/pinned.rs index 52b1aa46e2a..732c5d93dd7 100644 --- a/vortex-cuda/src/pinned.rs +++ b/vortex-cuda/src/pinned.rs @@ -309,6 +309,16 @@ impl PooledPinnedBuffer { .unwrap_or_else(|e| vortex_panic!("failed to access pinned host buffer: {e}")) } + #[cfg(target_os = "linux")] + pub(crate) fn truncate(&mut self, len: usize) { + let inner = self + .inner + .as_mut() + .unwrap_or_else(|| vortex_panic!("buffer already consumed")); + assert!(len <= inner.len()); + inner.set_logical_len(len); + } + /// Submits a non-blocking H2D DMA transfer and returns a device buffer. /// /// The pinned buffer is placed in the pool's inflight queue, gated on a `CudaEvent` marking diff --git a/vortex-cuda/src/pooled_read_at.rs b/vortex-cuda/src/pooled_read_at.rs index d3024140c96..7ee672f6574 100644 --- a/vortex-cuda/src/pooled_read_at.rs +++ b/vortex-cuda/src/pooled_read_at.rs @@ -3,6 +3,8 @@ use std::fs::File; use std::io; +#[cfg(target_os = "linux")] +use std::os::unix::fs::FileExt; use std::path::Path; use std::sync::Arc; @@ -15,6 +17,14 @@ use object_store::GetResultPayload; use object_store::ObjectStore; use object_store::ObjectStoreExt; use object_store::path::Path as ObjectPath; +#[cfg(target_os = "linux")] +use rustix::fs::AtFlags; +#[cfg(target_os = "linux")] +use rustix::fs::Mode; +#[cfg(target_os = "linux")] +use rustix::fs::OFlags; +#[cfg(target_os = "linux")] +use rustix::fs::StatxFlags; use vortex::array::buffer::BufferHandle; use vortex::buffer::Alignment; use vortex::buffer::ByteBuffer; @@ -35,6 +45,26 @@ pub const DEFAULT_FILE_CONCURRENCY: usize = 32; /// Default number of concurrent requests to allow for object store I/O. pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192; +#[cfg(target_os = "linux")] +const FALLBACK_DIRECT_IO_ALIGNMENT: usize = 4096; + +#[cfg(target_os = "linux")] +#[derive(Clone, Copy)] +struct DirectIoAlignment { + memory: usize, + offset: usize, +} + +#[cfg(target_os = "linux")] +impl Default for DirectIoAlignment { + fn default() -> Self { + Self { + memory: FALLBACK_DIRECT_IO_ALIGNMENT, + offset: FALLBACK_DIRECT_IO_ALIGNMENT, + } + } +} + /// File reader that uses CUDA pinned host memory for I/O buffers and transfers /// directly to the GPU. /// @@ -50,6 +80,9 @@ pub struct PooledFileReadAt { handle: Handle, pool: Arc, stream: VortexCudaStream, + direct_io: bool, + #[cfg(target_os = "linux")] + direct_io_alignment: DirectIoAlignment, } impl PooledFileReadAt { @@ -69,8 +102,56 @@ impl PooledFileReadAt { handle, pool, stream, + direct_io: false, + #[cfg(target_os = "linux")] + direct_io_alignment: DirectIoAlignment::default(), }) } + + /// Open a file for pooled direct I/O with direct device transfer. + /// + /// Direct I/O bypasses the operating system page cache. Unaligned logical reads are widened + /// to aligned physical reads and sliced to the requested range after transfer to the device. + #[cfg(target_os = "linux")] + pub fn open_direct( + path: impl AsRef, + handle: Handle, + pool: Arc, + stream: VortexCudaStream, + ) -> VortexResult { + let path = path.as_ref(); + let uri = Arc::from(path.to_string_lossy().to_string()); + let file = File::from( + rustix::fs::open( + path, + OFlags::RDONLY | OFlags::CLOEXEC | OFlags::DIRECT, + Mode::empty(), + ) + .map_err(io::Error::from)?, + ); + let direct_io_alignment = direct_io_alignment(&file); + let file = Arc::new(file); + Ok(Self { + uri, + file, + handle, + pool, + stream, + direct_io: true, + direct_io_alignment, + }) + } + + /// Return an error when direct I/O is requested on an unsupported platform. + #[cfg(not(target_os = "linux"))] + pub fn open_direct( + _path: impl AsRef, + _handle: Handle, + _pool: Arc, + _stream: VortexCudaStream, + ) -> VortexResult { + vortex::error::vortex_bail!("direct CUDA file I/O is only supported on Linux") + } } impl VortexReadAt for PooledFileReadAt { @@ -105,24 +186,166 @@ impl VortexReadAt for PooledFileReadAt { let handle = self.handle.clone(); let stream = self.stream.clone(); let pool = Arc::clone(&self.pool); + let direct_io = self.direct_io; + #[cfg(target_os = "linux")] + let direct_io_alignment = self.direct_io_alignment; async move { - let mut target = pool.get(length)?; + #[cfg(target_os = "linux")] + let (read_offset, read_length, requested_range) = if direct_io { + direct_io_range(offset, length, direct_io_alignment.offset)? + } else { + (offset, length, 0..length) + }; + #[cfg(not(target_os = "linux"))] + let (read_offset, read_length, requested_range) = { + vortex_ensure!( + !direct_io, + "direct CUDA file I/O is only supported on Linux" + ); + (offset, length, 0..length) + }; + #[cfg(target_os = "linux")] + let required_bytes = requested_range.end; + + let mut target = pool.get(read_length)?; let target = handle .spawn_blocking(move || { - read_exact_at(&file, target.as_mut_slice(), offset)?; + #[cfg(target_os = "linux")] + if direct_io { + let address = target.as_mut_slice().as_ptr() as usize; + if !address.is_multiple_of(direct_io_alignment.memory) { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!( + "pinned buffer address {address:#x} is not aligned to {} bytes", + direct_io_alignment.memory + ), + )); + } + + let bytes_read = read_direct_at( + &file, + target.as_mut_slice(), + read_offset, + required_bytes, + direct_io_alignment, + )?; + target.truncate(bytes_read); + } else { + read_exact_at(&file, target.as_mut_slice(), read_offset)?; + } + #[cfg(not(target_os = "linux"))] + read_exact_at(&file, target.as_mut_slice(), read_offset)?; Ok::<_, io::Error>(target) }) .await .map_err(VortexError::from)?; let cuda_buf = target.transfer_to_device(&stream)?; - Ok(BufferHandle::new_device(Arc::new(cuda_buf))) + Ok(BufferHandle::new_device(Arc::new(cuda_buf)).slice(requested_range)) } .boxed() } } +#[cfg(target_os = "linux")] +fn direct_io_range( + offset: u64, + length: usize, + alignment: usize, +) -> VortexResult<(u64, usize, std::ops::Range)> { + vortex_ensure!(alignment > 0, "direct I/O alignment must be non-zero"); + let alignment_u64 = u64::try_from(alignment)?; + let length_u64 = u64::try_from(length)?; + offset.checked_add(length_u64).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + let read_offset = offset / alignment_u64 * alignment_u64; + let prefix = usize::try_from(offset - read_offset)?; + let requested_end = prefix.checked_add(length).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + let read_length = requested_end + .checked_add(alignment - 1) + .ok_or_else(|| vortex_err!("direct I/O aligned length overflow"))? + / alignment + * alignment; + + Ok((read_offset, read_length, prefix..requested_end)) +} + +#[cfg(target_os = "linux")] +fn read_direct_at( + file: &File, + buffer: &mut [u8], + offset: u64, + required_bytes: usize, + alignment: DirectIoAlignment, +) -> io::Result { + let mut filled = 0; + while filled < required_bytes { + let filled_u64 = u64::try_from(filled) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let read_offset = offset + .checked_add(filled_u64) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let bytes_read = match file.read_at(&mut buffer[filled..], read_offset) { + Err(error) if error.kind() == io::ErrorKind::Interrupted => continue, + result => result?, + }; + if bytes_read == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned {filled} bytes, but {required_bytes} bytes were required" + ), + )); + } + filled += bytes_read; + if filled < required_bytes + && (!filled.is_multiple_of(alignment.offset) + || !filled.is_multiple_of(alignment.memory)) + { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned an unaligned short read of {filled} bytes before the required {required_bytes} bytes" + ), + )); + } + } + + Ok(filled) +} + +#[cfg(target_os = "linux")] +fn direct_io_alignment(file: &File) -> DirectIoAlignment { + let Ok(stat) = rustix::fs::statx( + file, + c"", + AtFlags::EMPTY_PATH | AtFlags::STATX_DONT_SYNC, + StatxFlags::DIOALIGN, + ) else { + return DirectIoAlignment::default(); + }; + if stat.stx_mask & StatxFlags::DIOALIGN.bits() == 0 { + return DirectIoAlignment::default(); + } + + let Ok(memory) = usize::try_from(stat.stx_dio_mem_align) else { + return DirectIoAlignment::default(); + }; + let Ok(offset) = usize::try_from(stat.stx_dio_offset_align) else { + return DirectIoAlignment::default(); + }; + if memory == 0 || offset == 0 { + return DirectIoAlignment::default(); + } + + DirectIoAlignment { memory, offset } +} + /// Object store reader that uses CUDA pinned host memory for I/O buffers and /// transfers directly to the GPU. /// @@ -350,3 +573,25 @@ impl VortexReadAt for PooledByteBufferReadAt { .boxed() } } + +#[cfg(test)] +mod tests { + #[cfg(target_os = "linux")] + use super::*; + + #[cfg(target_os = "linux")] + #[test] + fn widens_unaligned_direct_read_to_block_boundaries() -> VortexResult<()> { + assert_eq!(direct_io_range(5, 10, 4096)?, (0, 4096, 5..15)); + assert_eq!(direct_io_range(4090, 20, 4096)?, (0, 8192, 4090..4110)); + assert_eq!(direct_io_range(4096, 4096, 4096)?, (4096, 4096, 0..4096)); + assert_eq!(direct_io_range(513, 1, 512)?, (512, 512, 1..2)); + Ok(()) + } + + #[cfg(target_os = "linux")] + #[test] + fn rejects_overflowing_direct_read_range() { + assert!(direct_io_range(u64::MAX, 2, 4096).is_err()); + } +} From 6b265feb2d8c022cafa65103ff9d481189c9f97e Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Wed, 22 Jul 2026 17:01:36 +0100 Subject: [PATCH 2/8] feat[io]: add direct I/O for CPU file scans Signed-off-by: Onur Satici --- .github/workflows/sql-benchmarks.yml | 1 + Cargo.lock | 1 + benchmarks/datafusion-bench/src/lib.rs | 42 +++- benchmarks/datafusion-bench/src/main.rs | 2 +- vortex-cuda/ffi/cinclude/vortex_cuda.h | 11 +- vortex-datafusion/src/persistent/format.rs | 31 ++- vortex-file/src/open.rs | 32 ++- vortex-io/Cargo.toml | 3 + vortex-io/src/std_file/read_at.rs | 248 +++++++++++++++++++++ 9 files changed, 354 insertions(+), 17 deletions(-) diff --git a/.github/workflows/sql-benchmarks.yml b/.github/workflows/sql-benchmarks.yml index 7c4b69a2539..12b115fc6f5 100644 --- a/.github/workflows/sql-benchmarks.yml +++ b/.github/workflows/sql-benchmarks.yml @@ -532,6 +532,7 @@ jobs: env: VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + VORTEX_DIRECT_IO: "1" # Makes python output nicer COLUMNS: 120 strategy: diff --git a/Cargo.lock b/Cargo.lock index 33f17258b95..ccf297f64d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10304,6 +10304,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "rstest", + "rustix", "smol", "tempfile", "tokio", diff --git a/benchmarks/datafusion-bench/src/lib.rs b/benchmarks/datafusion-bench/src/lib.rs index c45aa99be2c..db465624de1 100644 --- a/benchmarks/datafusion-bench/src/lib.rs +++ b/benchmarks/datafusion-bench/src/lib.rs @@ -10,6 +10,7 @@ use datafusion::datasource::file_format::FileFormat; use datafusion::datasource::file_format::arrow::ArrowFormat; use datafusion::datasource::file_format::csv::CsvFormat; use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::provider::DefaultTableFactory; use datafusion::execution::SessionStateBuilder; use datafusion::execution::cache::DefaultListFilesCache; @@ -19,16 +20,22 @@ use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::prelude::SessionConfig; use datafusion::prelude::SessionContext; use datafusion_common::GetExt; +use datafusion_common::Result as DFResult; use object_store::ObjectStore; use object_store::aws::AmazonS3Builder; use object_store::gcp::GoogleCloudStorageBuilder; use object_store::local::LocalFileSystem; use url::Url; +use vortex::array::memory::MemorySessionExt; +use vortex::io::VortexReadAt; +use vortex::io::session::RuntimeSessionExt; +use vortex::io::std_file::FileReadAt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_datafusion::VortexFormat; use vortex_datafusion::VortexFormatFactory; use vortex_datafusion::VortexTableOptions; +use vortex_datafusion::reader::VortexReaderFactory; #[expect(clippy::expect_used)] pub fn get_session_context() -> SessionContext { @@ -106,20 +113,47 @@ pub fn make_object_store( } } -pub fn format_to_df_format(format: Format) -> Arc { +pub fn format_to_df_format(format: Format, source: &Url) -> Arc { match format { Format::Csv => Arc::new(CsvFormat::default()) as _, Format::Arrow => Arc::new(ArrowFormat), Format::Parquet => Arc::new(ParquetFormat::new()), - Format::OnDiskVortex | Format::VortexCompact | Format::VortexNative => Arc::new( - VortexFormat::new_with_options(SESSION.clone(), vortex_table_options()), - ), + Format::OnDiskVortex | Format::VortexCompact | Format::VortexNative => { + let format = VortexFormat::new_with_options(SESSION.clone(), vortex_table_options()); + let format = if source.scheme() == "file" + && std::env::var("VORTEX_DIRECT_IO").is_ok_and(|value| value == "1") + { + format.with_vortex_reader_factory(Arc::new(DirectIoVortexReaderFactory)) + } else { + format + }; + Arc::new(format) + } Format::OnDiskDuckDB | Format::Lance => { unimplemented!("Format {format} cannot be turned into a DataFusion `FileFormat`") } } } +#[derive(Debug)] +struct DirectIoVortexReaderFactory; + +impl VortexReaderFactory for DirectIoVortexReaderFactory { + fn create_reader( + &self, + file: &PartitionedFile, + session: &vortex::session::VortexSession, + ) -> DFResult> { + let path = LocalFileSystem::default() + .path_to_filesystem(file.path()) + .map_err(|error| datafusion_common::DataFusionError::External(Box::new(error)))?; + let reader = + FileReadAt::open_direct_with_allocator(path, session.handle(), session.allocator()) + .map_err(|error| datafusion_common::DataFusionError::External(Box::new(error)))?; + Ok(Arc::new(reader)) + } +} + fn vortex_table_options() -> VortexTableOptions { let mut opts = VortexTableOptions::default(); diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 67521a8147e..84a709e7ccc 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -263,7 +263,7 @@ async fn register_benchmark_tables( } _ => { let benchmark_base = benchmark.data_url().join(&format!("{}/", format.name()))?; - let file_format = format_to_df_format(format); + let file_format = format_to_df_format(format, benchmark.data_url()); for table in benchmark.table_specs().iter() { let pattern = benchmark.pattern(table.name, format); diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 6f1e171c739..09f5261cd7a 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -145,12 +145,11 @@ int vx_cuda_scan_path_arrow_device_stream_batch_rows(const vx_session *session, * `vx_cuda_scan_path_arrow_device_stream`. Pass NULL or a zero-initialized options struct to use * buffered file I/O and layout-derived batch splitting. */ -int vx_cuda_scan_path_arrow_device_stream_with_options( - const vx_session *session, - vx_view path, - const vx_cuda_scan_options *options, - struct ArrowDeviceArrayStream *out_stream, - vx_error **error_out); +int vx_cuda_scan_path_arrow_device_stream_with_options(const vx_session *session, + vx_view path, + const vx_cuda_scan_options *options, + struct ArrowDeviceArrayStream *out_stream, + vx_error **error_out); /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 779b74b212e..044e401a426 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -65,6 +65,7 @@ use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; use super::cache::CachedVortexMetadata; +use super::reader::VortexReaderFactory; use super::sink::VortexSink; use super::source::VortexSource; use crate::PrecisionExt as _; @@ -123,6 +124,7 @@ const DEFAULT_FOOTER_INITIAL_READ_SIZE_BYTES: usize = MAX_POSTSCRIPT_SIZE as usi pub struct VortexFormat { session: VortexSession, opts: VortexTableOptions, + vortex_reader_factory: Option>, } impl Debug for VortexFormat { @@ -398,7 +400,11 @@ impl VortexFormat { /// Creates a format with explicit [`VortexTableOptions`]. pub fn new_with_options(session: VortexSession, opts: VortexTableOptions) -> Self { - Self { session, opts } + Self { + session, + opts, + vortex_reader_factory: None, + } } /// Returns the format-specific configuration that will be copied into the @@ -406,6 +412,20 @@ impl VortexFormat { pub fn options(&self) -> &VortexTableOptions { &self.opts } + + /// Sets a custom factory for the underlying [`VortexReadAt`] used by scans. + /// + /// Metadata inference continues to use the DataFusion object store. The custom factory is + /// applied to the [`VortexSource`] created for physical scan execution. + /// + /// [`VortexReadAt`]: vortex::io::VortexReadAt + pub fn with_vortex_reader_factory( + mut self, + vortex_reader_factory: Arc, + ) -> Self { + self.vortex_reader_factory = Some(vortex_reader_factory); + self + } } #[async_trait] @@ -702,9 +722,12 @@ impl FileFormat for VortexFormat { } fn file_source(&self, table_schema: TableSchema) -> Arc { - Arc::new( - VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()), - ) as _ + let mut source = + VortexSource::new(table_schema, self.session.clone()).with_options(self.opts.clone()); + if let Some(vortex_reader_factory) = &self.vortex_reader_factory { + source = source.with_vortex_reader_factory(Arc::clone(vortex_reader_factory)); + } + Arc::new(source) as _ } } diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index 79e8e671240..f99a431afce 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -75,6 +75,8 @@ pub struct VortexOpenOptions { labels: Vec