diff --git a/.github/workflows/pr-bench-runner.yml b/.github/workflows/pr-bench-runner.yml index 017c5468772..d4791638adf 100644 --- a/.github/workflows/pr-bench-runner.yml +++ b/.github/workflows/pr-bench-runner.yml @@ -73,6 +73,11 @@ jobs: needs: build if: ${{ !cancelled() && needs.build.result == 'success' }} timeout-minutes: 120 + env: + # BENCHMARK EXPERIMENT -- revert before merging. + # Each benchmark writes its own Vortex files and reads them back in this job. + VORTEX_DIRECT_IO: "1" + VORTEX_SEGMENT_PADDING: "always" runs-on: >- ${{ github.repository == 'vortex-data/vortex' && format('runs-on={0}/runner=bench-dedicated/family=c8gd.metal-24xl/image=ubuntu24-full-arm64-pre-v2/tag={1}{2}', github.run_id, inputs.benchmark_id, github.event.pull_request.head.repo.fork == false && '/extras=s3-cache' || '') diff --git a/.github/workflows/sql-bench-matrix.yml b/.github/workflows/sql-bench-matrix.yml index d745eae125b..ac6f33b3460 100644 --- a/.github/workflows/sql-bench-matrix.yml +++ b/.github/workflows/sql-bench-matrix.yml @@ -103,6 +103,11 @@ jobs: FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" # Makes python output nicer COLUMNS: 120 + # BENCHMARK EXPERIMENT -- revert before merging. + # `prepare-data` and the query runs share this job, so these cover both writing the + # Vortex files and reading them back. + VORTEX_DIRECT_IO: "1" + VORTEX_SEGMENT_PADDING: "always" strategy: fail-fast: false matrix: diff --git a/Cargo.lock b/Cargo.lock index 6647a4ff108..d79aeeb4749 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11054,6 +11054,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "rstest", + "tempfile", "tokio", "tracing", "url", @@ -11160,6 +11161,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "rstest", + "rustix", "smol", "tempfile", "tokio", diff --git a/vortex-bench/Cargo.toml b/vortex-bench/Cargo.toml index a63f954a637..b294c09c6c3 100644 --- a/vortex-bench/Cargo.toml +++ b/vortex-bench/Cargo.toml @@ -17,6 +17,10 @@ version = { workspace = true } name = "data-gen" test = false +[[bin]] +name = "direct-io" +test = false + [lints] workspace = true diff --git a/vortex-bench/src/bin/direct-io.rs b/vortex-bench/src/bin/direct-io.rs new file mode 100644 index 00000000000..2bdcb22126b --- /dev/null +++ b/vortex-bench/src/bin/direct-io.rs @@ -0,0 +1,619 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measure direct I/O reads and block-aligned segment writes. +//! +//! `convert` writes Vortex files under each [`SegmentPadding`] policy, `analyze` replays every +//! policy against an already-written file's segment map to price it without rewriting, and `scan` +//! times full scans through the layout reader's `ScanBuilder` with buffered and direct reads. + +use std::collections::BTreeMap; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use clap::Parser; +use clap::Subcommand; +use clap::ValueEnum; +use futures::StreamExt; +use futures::TryStreamExt; +use humansize::DECIMAL; +use humansize::format_size; +use parquet::arrow::ParquetRecordBatchStreamBuilder; +use vortex::array::ArrayRef; +use vortex::array::memory::BufferAllocatorRef; +use vortex::array::stream::ArrayStreamAdapter; +use vortex::array::stream::ArrayStreamExt; +use vortex::buffer::Alignment; +use vortex::dtype::FieldName; +use vortex::expr::root; +use vortex::expr::select; +use vortex::file::OpenOptionsSessionExt; +use vortex::file::SegmentSpec; +use vortex::file::VortexFile; +use vortex::file::WriteOptionsSessionExt; +use vortex::file::segments::SegmentPadding; +use vortex::io::runtime::Handle; +use vortex::io::session::RuntimeSessionExt; +use vortex::io::std_file::DEFAULT_CONCURRENCY; +use vortex::io::std_file::FileReadAt; +use vortex::io::std_file::FileReadAtOptions; +use vortex::layout::segments::SegmentId; +use vortex_arrow::ArrowSessionExt; +use vortex_bench::SESSION; +use vortex_bench::conversions::parquet_to_vortex_stream; +use vortex_bench::setup_logging_and_tracing; + +/// Block size every policy in this tool aligns to. +const BLOCK: u64 = 4096; + +/// The same block size as an [`Alignment`], for building padding policies. +const BLOCK_ALIGNMENT: Alignment = Alignment::new(4096); + +#[derive(Parser)] +#[command( + name = "direct-io", + about = "Direct I/O read and segment alignment measurements" +)] +struct Args { + #[arg(short, long, global = true)] + verbose: bool, + #[command(subcommand)] + command: Command, +} + +#[derive(Subcommand)] +enum Command { + /// Convert Parquet files to Vortex under each padding policy. + Convert { + /// Parquet files, or directories of them. + #[arg(required = true)] + inputs: Vec, + /// Directory to write Vortex files into. + #[arg(short, long)] + out: PathBuf, + /// Padding policies to write. + #[arg(long, value_delimiter = ',', default_values = ["none", "block"])] + padding: Vec, + }, + /// Price every padding policy against already-written Vortex files. + Analyze { + /// Vortex files, or directories of them. + #[arg(required = true)] + inputs: Vec, + /// Overhead ratios to sweep for the proportional policy. + #[arg(long, value_delimiter = ',', default_values_t = [4u32, 8, 16, 32, 64, 128, 256])] + ratios: Vec, + }, + /// Time full scans with buffered and direct reads. + Scan { + /// Vortex files, or directories of them. + #[arg(required = true)] + inputs: Vec, + #[arg(short, long, default_value_t = 3)] + iterations: usize, + /// Restrict the scan to these top-level columns. + #[arg(long, value_delimiter = ',')] + columns: Vec, + /// Drop the page cache before every iteration, measuring cold reads. + #[arg(long)] + cold: bool, + /// Resolve every segment without decoding, isolating the read pipeline from decompression. + #[arg(long)] + io_only: bool, + /// Number of files to scan concurrently. + #[arg(long, default_value_t = 4)] + concurrency: usize, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, ValueEnum)] +enum Padding { + /// Pack segments contiguously. + None, + /// Start every segment on a 4KiB boundary. + Block, + /// Pack consecutive segments into shared 4KiB blocks. + Grouped, + /// Block-align a segment when the padding is within 1/64 of its length. + Proportional, +} + +impl Padding { + fn policy(self) -> SegmentPadding { + match self { + Self::None => SegmentPadding::None, + Self::Block => SegmentPadding::block_aligned(), + Self::Grouped => SegmentPadding::grouped(), + Self::Proportional => SegmentPadding::proportional(), + } + } + + fn suffix(self) -> &'static str { + match self { + Self::None => "none", + Self::Block => "block", + Self::Grouped => "grouped", + Self::Proportional => "proportional", + } + } +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let args = Args::parse(); + setup_logging_and_tracing(args.verbose, false)?; + + match args.command { + Command::Convert { + inputs, + out, + padding, + } => convert(&collect(&inputs, "parquet")?, &out, &padding).await, + Command::Analyze { inputs, ratios } => analyze(&collect(&inputs, "vortex")?, &ratios).await, + Command::Scan { + inputs, + iterations, + columns, + cold, + io_only, + concurrency, + } => { + scan( + &collect(&inputs, "vortex")?, + iterations, + ScanConfig { + columns, + cold, + io_only, + concurrency, + }, + ) + .await + } + } +} + +/// Expand directories into the files inside them with the given extension. +fn collect(inputs: &[PathBuf], extension: &str) -> anyhow::Result> { + let mut files = Vec::new(); + for input in inputs { + if input.is_dir() { + for entry in fs::read_dir(input)? { + let path = entry?.path(); + if path.extension().is_some_and(|ext| ext == extension) { + files.push(path); + } + } + } else { + files.push(input.clone()); + } + } + files.sort(); + anyhow::ensure!(!files.is_empty(), "no .{extension} files found"); + Ok(files) +} + +async fn convert(inputs: &[PathBuf], out: &Path, padding: &[Padding]) -> anyhow::Result<()> { + fs::create_dir_all(out)?; + + let mut totals: BTreeMap<&str, u64> = BTreeMap::new(); + let mut parquet_total = 0; + for input in inputs { + let stem = input + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("data"); + parquet_total += fs::metadata(input)?.len(); + + for padding in padding { + let output = out.join(format!("{stem}.{}.vortex", padding.suffix())); + let elapsed = if output.exists() { + Duration::ZERO + } else { + let start = Instant::now(); + write_vortex(input, &output, padding.policy()).await?; + start.elapsed() + }; + let size = fs::metadata(&output)?.len(); + *totals.entry(padding.suffix()).or_default() += size; + println!( + "{stem:<24} {:<13} {:>12} ({:.1}s)", + padding.suffix(), + format_size(size, DECIMAL), + elapsed.as_secs_f64(), + ); + } + } + + println!("\nparquet total: {}", format_size(parquet_total, DECIMAL)); + let baseline = totals.get("none").copied(); + for (policy, size) in &totals { + let growth = baseline + .map(|baseline| format!(" {:+.3}%", percent_change(baseline, *size))) + .unwrap_or_default(); + println!("{policy:<13} {:>12}{growth}", format_size(*size, DECIMAL)); + } + Ok(()) +} + +async fn write_vortex(input: &Path, output: &Path, padding: SegmentPadding) -> anyhow::Result<()> { + let file = tokio::fs::File::open(input).await?; + let builder = ParquetRecordBatchStreamBuilder::new(file).await?; + let dtype = SESSION + .arrow() + .from_arrow_schema(builder.schema().as_ref())?; + let stream = parquet_to_vortex_stream(builder.build()?); + + let mut out = tokio::fs::File::create(output).await?; + SESSION + .write_options() + .with_segment_padding(padding) + .write( + &mut out, + ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)), + ) + .await?; + Ok(()) +} + +/// Coalescing window the file reader applies to local files, mirroring `CoalesceConfig::file`. +const COALESCE_DISTANCE: u64 = 1 << 20; +const COALESCE_MAX_SIZE: u64 = 4 << 20; + +/// What a padding policy would cost, replayed over an existing file's segment map. +#[derive(Default)] +struct Replay { + file_size: u64, + padding: u64, + aligned_segments: usize, + /// Segment bytes living in a block-aligned segment. + aligned_bytes: u64, + /// Bytes a direct-I/O reader transfers when each segment is read on its own. + uncoalesced_bytes: u64, + /// Bytes a direct-I/O reader transfers once neighbouring segments are coalesced. + coalesced_bytes: u64, +} + +fn replay(specs: &[SegmentSpec], start: u64, policy: SegmentPadding) -> Replay { + let mut offset = start; + let mut replay = Replay::default(); + // Placed segments, in offset order, so the coalescing pass can walk them. + let mut placed: Vec<(u64, u64)> = Vec::with_capacity(specs.len()); + + for spec in specs { + let length = u64::from(spec.length); + let pad = policy.padding(offset, length, spec.alignment); + offset += pad; + replay.padding += pad; + if offset.is_multiple_of(BLOCK) { + replay.aligned_segments += 1; + replay.aligned_bytes += length; + } + // A direct read is widened to the blocks it overlaps, so a segment straddling a boundary + // transfers one more block than its length alone implies. + replay.uncoalesced_bytes += widen(offset, length); + placed.push((offset, length)); + offset += length; + } + replay.file_size = offset; + + // Coalesce neighbours the way the read driver does, then widen each physical read once. + let mut run: Option<(u64, u64)> = None; + for (offset, length) in placed { + run = Some(match run { + Some((start, end)) + if offset.saturating_sub(end) <= COALESCE_DISTANCE + && offset + length - start <= COALESCE_MAX_SIZE => + { + (start, end.max(offset + length)) + } + Some((start, end)) => { + replay.coalesced_bytes += widen(start, end - start); + (offset, offset + length) + } + None => (offset, offset + length), + }); + } + if let Some((start, end)) = run { + replay.coalesced_bytes += widen(start, end - start); + } + + replay +} + +/// Bytes transferred by a direct read of `offset..offset + length`, widened to whole blocks. +fn widen(offset: u64, length: u64) -> u64 { + if length == 0 { + return 0; + } + (offset % BLOCK + length).div_ceil(BLOCK) * BLOCK +} + +async fn analyze(inputs: &[PathBuf], ratios: &[u32]) -> anyhow::Result<()> { + let mut all_specs = Vec::new(); + let mut baseline_size = 0; + let mut trailer = 0; + + for input in inputs { + let file = SESSION.open_options().open_path(input).await?; + let specs = file.footer().segment_map().to_vec(); + let size = fs::metadata(input)?.len(); + let segment_end = specs + .iter() + .map(|spec| spec.offset + u64::from(spec.length)) + .max() + .unwrap_or(0); + trailer += size - segment_end; + baseline_size += size; + all_specs.push(specs); + } + + let segments: Vec = all_specs.iter().flatten().copied().collect(); + let data: u64 = segments.iter().map(|s| u64::from(s.length)).sum(); + println!("files: {}", inputs.len()); + println!("segments: {}", segments.len()); + println!("segment bytes: {}", format_size(data, DECIMAL)); + println!("file bytes: {}", format_size(baseline_size, DECIMAL)); + print_size_distribution(&segments); + + println!( + "\n{:<28} {:>13} {:>8} {:>9} {:>9} {:>13} {:>13}", + "policy", "file size", "growth", "aligned", "of bytes", "1 seg/io", "coalesced" + ); + // The read columns are deltas against the packed baseline, so it must be replayed first. + let mut policies = vec![ + ("none".to_string(), SegmentPadding::None), + ("always (4KiB)".to_string(), SegmentPadding::block_aligned()), + ("grouped (4KiB)".to_string(), SegmentPadding::grouped()), + ]; + for ratio in ratios { + policies.push(( + format!("proportional 1/{ratio}"), + SegmentPadding::Proportional { + block: BLOCK_ALIGNMENT, + max_overhead_ratio: *ratio, + }, + )); + } + + // Every file starts its segments after the magic bytes, so replay each separately and sum. + let totals: Vec<(String, Replay)> = policies + .into_iter() + .map(|(name, policy)| { + let mut total = Replay::default(); + for specs in &all_specs { + let start = specs.first().map(|spec| spec.offset).unwrap_or(0); + let replayed = replay(specs, start, policy); + total.file_size += replayed.file_size; + total.padding += replayed.padding; + total.aligned_segments += replayed.aligned_segments; + total.aligned_bytes += replayed.aligned_bytes; + total.uncoalesced_bytes += replayed.uncoalesced_bytes; + total.coalesced_bytes += replayed.coalesced_bytes; + } + (name, total) + }) + .collect(); + + // Price every policy against the contiguously packed layout it would replace. + let base = &totals[0].1; + for (name, total) in &totals { + println!( + "{name:<28} {:>13} {:>7.3}% {:>8.1}% {:>8.1}% {:>7.2}% {:>7.2}%", + format_size(total.file_size + trailer, DECIMAL), + percent_change(baseline_size, total.file_size + trailer), + 100.0 * total.aligned_segments as f64 / segments.len() as f64, + 100.0 * total.aligned_bytes as f64 / data as f64, + percent_change(base.uncoalesced_bytes, total.uncoalesced_bytes), + percent_change(base.coalesced_bytes, total.coalesced_bytes), + ); + } + println!( + "\nbytes read, packed baseline: {} (1 seg/io), {} (coalesced), for {} of segments", + format_size(base.uncoalesced_bytes, DECIMAL), + format_size(base.coalesced_bytes, DECIMAL), + format_size(data, DECIMAL), + ); + println!( + "\n\"read\" columns are the bytes a direct-I/O reader transfers to read every segment,\n\ + widened to {BLOCK}-byte blocks. The coalesced column applies the reader's own\n\ + {}/{} coalescing window first.", + format_size(COALESCE_DISTANCE, DECIMAL), + format_size(COALESCE_MAX_SIZE, DECIMAL), + ); + Ok(()) +} + +fn print_size_distribution(segments: &[SegmentSpec]) { + let mut lengths: Vec = segments.iter().map(|s| u64::from(s.length)).collect(); + lengths.sort_unstable(); + let quantile = |numerator: usize, denominator: usize| { + lengths[(lengths.len() * numerator / denominator).min(lengths.len() - 1)] + }; + println!( + "segment size: p50={} p90={} p99={} max={}", + format_size(quantile(1, 2), DECIMAL), + format_size(quantile(9, 10), DECIMAL), + format_size(quantile(99, 100), DECIMAL), + format_size(*lengths.last().unwrap_or(&0), DECIMAL), + ); + for threshold in [BLOCK, 4 * BLOCK, 16 * BLOCK] { + let count = lengths.partition_point(|len| *len < threshold); + let bytes: u64 = lengths.iter().take_while(|len| **len < threshold).sum(); + println!( + " < {:>7}: {:>7} segments ({:>4.1}%), {:>10} ({:.1}% of bytes)", + format_size(threshold, DECIMAL), + count, + 100.0 * count as f64 / lengths.len() as f64, + format_size(bytes, DECIMAL), + 100.0 * bytes as f64 / lengths.iter().sum::().max(1) as f64, + ); + } +} + +fn percent_change(baseline: u64, value: u64) -> f64 { + 100.0 * (value as f64 - baseline as f64) / baseline as f64 +} + +struct ScanConfig { + columns: Vec, + cold: bool, + io_only: bool, + concurrency: usize, +} + +async fn scan(inputs: &[PathBuf], iterations: usize, config: ScanConfig) -> anyhow::Result<()> { + let bytes: u64 = inputs + .iter() + .map(|input| fs::metadata(input).map(|meta| meta.len())) + .sum::>()?; + + println!( + "{} file(s), {}, {iterations} iteration(s), {} cache, {}", + inputs.len(), + format_size(bytes, DECIMAL), + if config.cold { "cold" } else { "warm" }, + if config.io_only { + "segments only" + } else { + "full scan" + }, + ); + println!( + "\n{:<10} {:>10} {:>10} {:>10} {:>12}", + "mode", "min", "median", "max", "throughput" + ); + + let mut baseline = None; + for direct in [false, true] { + let options = read_options(direct); + let mut timings = Vec::with_capacity(iterations); + let mut rows = 0; + for _ in 0..iterations { + if config.cold { + drop_caches()?; + } + let start = Instant::now(); + rows = scan_once(inputs, &config, options).await?; + timings.push(start.elapsed()); + } + timings.sort(); + let median = timings[timings.len() / 2]; + let throughput = bytes as f64 / median.as_secs_f64() / 1e6; + let label = if direct { "direct" } else { "buffered" }; + let delta = baseline + .map(|base: Duration| { + format!( + " {:+.1}%", + 100.0 * (median.as_secs_f64() - base.as_secs_f64()) / base.as_secs_f64() + ) + }) + .unwrap_or_default(); + println!( + "{label:<10} {:>9.3}s {:>9.3}s {:>9.3}s {throughput:>9.0} MB/s{delta}", + timings[0].as_secs_f64(), + median.as_secs_f64(), + timings[timings.len() - 1].as_secs_f64(), + ); + baseline.get_or_insert(median); + anyhow::ensure!(rows > 0, "scan produced no rows"); + } + Ok(()) +} + +fn read_options(direct: bool) -> FileReadAtOptions { + #[cfg(target_os = "linux")] + if direct { + return FileReadAtOptions::default().with_direct_io(); + } + let _ = direct; + FileReadAtOptions::default() +} + +async fn scan_once( + inputs: &[PathBuf], + config: &ScanConfig, + options: FileReadAtOptions, +) -> anyhow::Result { + let handle = SESSION.handle(); + let scans = inputs.iter().map(|input| { + let handle = handle.clone(); + async move { scan_file(input, config, options, handle).await } + }); + let rows: Vec = futures::stream::iter(scans) + .buffer_unordered(config.concurrency.max(1)) + .try_collect() + .await?; + Ok(rows.into_iter().sum()) +} + +async fn scan_file( + input: &Path, + config: &ScanConfig, + options: FileReadAtOptions, + handle: Handle, +) -> anyhow::Result { + let reader = FileReadAt::open_with_options( + input, + handle, + BufferAllocatorRef::statically_allocated(), + options, + )?; + let file: VortexFile = SESSION.open_options().open(Arc::new(reader)).await?; + + if config.io_only { + return read_all_segments(&file).await; + } + + let mut scan = file.scan()?; + if !config.columns.is_empty() { + let names: Vec = config + .columns + .iter() + .map(|name| FieldName::from(name.as_str())) + .collect(); + let projection = select(names, root()) + .optimize_recursive(file.dtype())? + .bind(file.dtype())?; + scan = scan.with_projection(projection); + } + + let mut stream = Box::pin(scan.into_array_stream()?); + let mut rows = 0; + while let Some(array) = stream.next().await { + let array: ArrayRef = array?; + rows += array.len(); + } + Ok(rows) +} + +/// Pull every segment through the file's read pipeline without decoding any of them. +/// +/// This exercises the same request registration, coalescing, and concurrency limiting that a scan +/// drives, so it isolates the cost of the reads themselves from decompression. +async fn read_all_segments(file: &VortexFile) -> anyhow::Result { + let source = file.segment_source(); + let segment_count = file.footer().segment_map().len(); + let requests = (0..u32::try_from(segment_count)?).map(|id| { + let request = source.request(SegmentId::from(id)); + async move { request.await.map(|buffer| buffer.len()) } + }); + let lengths: Vec = futures::stream::iter(requests) + .buffered(DEFAULT_CONCURRENCY) + .try_collect() + .await?; + Ok(lengths.into_iter().sum()) +} + +/// Evict the page cache so the next scan reads from the device. +fn drop_caches() -> anyhow::Result<()> { + use std::io::Write; + + fs::File::create("/proc/sys/vm/drop_caches") + .and_then(|mut file| file.write_all(b"3")) + .map_err(|e| anyhow::anyhow!("cannot drop the page cache (needs root): {e}")) +} diff --git a/vortex-cuda/src/pooled_read_at/file/direct.rs b/vortex-cuda/src/pooled_read_at/file/direct.rs index 53366c36915..96aaed60406 100644 --- a/vortex-cuda/src/pooled_read_at/file/direct.rs +++ b/vortex-cuda/src/pooled_read_at/file/direct.rs @@ -2,54 +2,18 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::fs::File; -use std::io; -use std::ops::Range; -use std::os::unix::fs::FileExt; use std::path::Path; use std::sync::Arc; -use rustix::fs::AtFlags; -use rustix::fs::Mode; -use rustix::fs::OFlags; -use rustix::fs::StatxFlags; use vortex::error::VortexResult; use vortex::error::vortex_ensure; -use vortex::error::vortex_err; +use vortex::io::std_file::DirectIoConstraints; +use vortex::io::std_file::open_direct; use super::FileReadBackend; use super::PooledHostRead; use crate::pinned::PinnedByteBufferPool; -/// Conservative direct-I/O alignment used when Linux cannot report the filesystem constraints. -/// -/// A page-sized fallback is accepted by common block devices and filesystems. If the actual -/// requirement is stricter, the read fails with the underlying `EINVAL`. -const FALLBACK_DIRECT_IO_ALIGNMENT: usize = 4096; - -#[derive(Clone, Copy)] -struct DirectIoConstraints { - /// Required alignment of the address of the userspace I/O buffer. - memory_alignment: usize, - /// Required alignment of both the file offset and the I/O length. - offset_alignment: usize, -} - -impl Default for DirectIoConstraints { - fn default() -> Self { - Self { - memory_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, - offset_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, - } - } -} - -#[derive(Debug, PartialEq, Eq)] -struct DirectIoRange { - read_offset: u64, - read_length: usize, - requested_range: Range, -} - pub(super) struct DirectFileReadBackend { file: File, constraints: DirectIoConstraints, @@ -57,15 +21,8 @@ pub(super) struct DirectFileReadBackend { impl DirectFileReadBackend { pub(super) fn open(path: &Path) -> VortexResult { - let file = File::from( - rustix::fs::open( - path, - OFlags::RDONLY | OFlags::CLOEXEC | OFlags::DIRECT, - Mode::empty(), - ) - .map_err(io::Error::from)?, - ); - let constraints = direct_io_constraints(&file)?; + let file = open_direct(path)?; + let constraints = DirectIoConstraints::probe(&file)?; Ok(Self { file, constraints }) } } @@ -81,21 +38,20 @@ impl FileReadBackend for DirectFileReadBackend { offset: u64, length: usize, ) -> VortexResult { - let direct_range = direct_io_range(offset, length, self.constraints.offset_alignment)?; + let direct_range = self.constraints.widen(offset, length)?; let mut buffer = pool.get(direct_range.read_length)?; let address = buffer.as_mut_slice().as_ptr() as usize; vortex_ensure!( - address.is_multiple_of(self.constraints.memory_alignment), + address.is_multiple_of(self.constraints.memory_alignment()), "pinned buffer address {address:#x} is not aligned to {} bytes", - self.constraints.memory_alignment + self.constraints.memory_alignment() ); - let bytes_read = read_direct_at( + let bytes_read = self.constraints.read_at( &self.file, buffer.as_mut_slice(), direct_range.read_offset, direct_range.requested_range.end, - self.constraints, )?; buffer.truncate(bytes_read); Ok(PooledHostRead { @@ -104,185 +60,3 @@ impl FileReadBackend for DirectFileReadBackend { }) } } - -fn direct_io_range(offset: u64, length: usize, alignment: usize) -> VortexResult { - vortex_ensure!(alignment > 0, "direct I/O alignment must be non-zero"); - if length == 0 { - return Ok(DirectIoRange { - read_offset: offset, - read_length: 0, - requested_range: 0..0, - }); - } - - let alignment_u64 = u64::try_from(alignment)?; - let length_u64 = u64::try_from(length)?; - let requested_end = offset.checked_add(length_u64).ok_or_else(|| { - vortex_err!("direct I/O range overflow: offset={offset}, length={length}") - })?; - let read_offset = offset - offset % alignment_u64; - let read_end = requested_end - .checked_next_multiple_of(alignment_u64) - .ok_or_else(|| vortex_err!("direct I/O aligned end overflow"))?; - let read_length = usize::try_from(read_end - read_offset)?; - let slice_start = usize::try_from(offset - read_offset)?; - let slice_end = slice_start.checked_add(length).ok_or_else(|| { - vortex_err!("direct I/O range overflow: offset={offset}, length={length}") - })?; - - Ok(DirectIoRange { - read_offset, - read_length, - requested_range: slice_start..slice_end, - }) -} - -fn read_direct_at( - file: &File, - buffer: &mut [u8], - offset: u64, - required_bytes: usize, - constraints: DirectIoConstraints, -) -> io::Result { - let mut initialized = 0; - while initialized < required_bytes { - let initialized_u64 = u64::try_from(initialized) - .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; - let read_offset = offset - .checked_add(initialized_u64) - .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; - let bytes_read = match file.read_at(&mut buffer[initialized..], 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 {initialized} bytes, but {required_bytes} bytes were required" - ), - )); - } - initialized += bytes_read; - if initialized < required_bytes - && (!initialized.is_multiple_of(constraints.offset_alignment) - || !initialized.is_multiple_of(constraints.memory_alignment)) - { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - format!( - "direct read returned an unaligned short read of {initialized} bytes before the required {required_bytes} bytes" - ), - )); - } - } - - Ok(initialized) -} - -fn direct_io_constraints(file: &File) -> VortexResult { - let Ok(stat) = rustix::fs::statx( - file, - c"", - AtFlags::EMPTY_PATH | AtFlags::STATX_DONT_SYNC, - StatxFlags::DIOALIGN, - ) else { - return Ok(DirectIoConstraints::default()); - }; - if stat.stx_mask & StatxFlags::DIOALIGN.bits() == 0 { - return Ok(DirectIoConstraints::default()); - } - - let Ok(memory_alignment) = usize::try_from(stat.stx_dio_mem_align) else { - return Ok(DirectIoConstraints::default()); - }; - let Ok(offset_alignment) = usize::try_from(stat.stx_dio_offset_align) else { - return Ok(DirectIoConstraints::default()); - }; - if memory_alignment == 0 || offset_alignment == 0 { - return Ok(DirectIoConstraints::default()); - } - vortex_ensure!( - memory_alignment.is_power_of_two(), - "direct I/O memory alignment must be a power of two, got {memory_alignment}" - ); - vortex_ensure!( - offset_alignment.is_power_of_two(), - "direct I/O offset alignment must be a power of two, got {offset_alignment}" - ); - - Ok(DirectIoConstraints { - memory_alignment, - offset_alignment, - }) -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - - use super::*; - - #[rstest] - #[case(0, 0, 4096, 0, 0, 0)] - #[case(5, 0, 4096, 5, 0, 0)] - #[case(5, 10, 4096, 0, 4096, 5)] - #[case(4090, 20, 4096, 0, 8192, 4090)] - #[case(4096, 4096, 4096, 4096, 4096, 0)] - #[case(513, 1, 512, 512, 512, 1)] - #[case(4096, 8193, 4096, 4096, 12288, 0)] - fn widens_direct_read_to_block_boundaries( - #[case] offset: u64, - #[case] length: usize, - #[case] alignment: usize, - #[case] expected_offset: u64, - #[case] expected_length: usize, - #[case] expected_prefix: usize, - ) -> VortexResult<()> { - assert_eq!( - direct_io_range(offset, length, alignment)?, - DirectIoRange { - read_offset: expected_offset, - read_length: expected_length, - requested_range: expected_prefix..expected_prefix + length, - } - ); - Ok(()) - } - - #[rstest] - #[case(u64::MAX, 2, 4096)] - #[case(0, 1, 0)] - fn rejects_invalid_direct_read_range( - #[case] offset: u64, - #[case] length: usize, - #[case] alignment: usize, - ) { - assert!(direct_io_range(offset, length, alignment).is_err()); - } - - #[test] - fn aligned_ranges_cover_requested_bytes() -> VortexResult<()> { - for alignment in [512, 4096] { - for offset in 0..alignment * 2 { - for length in [0, 1, alignment - 1, alignment, alignment + 1] { - let range = direct_io_range(offset as u64, length, alignment)?; - if length == 0 { - assert_eq!(range.read_length, 0); - continue; - } - - assert_eq!(range.read_offset % alignment as u64, 0); - assert_eq!(range.read_length % alignment, 0); - assert_eq!(range.requested_range.len(), length); - assert!(range.requested_range.end <= range.read_length); - assert_eq!( - range.read_offset + range.requested_range.start as u64, - offset as u64 - ); - } - } - } - Ok(()) - } -} diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 78ff4d1d5b6..8b8ccae5654 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -64,6 +64,7 @@ vortex-zstd = { workspace = true, optional = true } allocator-api2 = { workspace = true } divan = { workspace = true } rstest = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-io = { workspace = true, features = ["tokio"] } diff --git a/vortex-file/src/segments/mod.rs b/vortex-file/src/segments/mod.rs index d96620fca23..beda775614e 100644 --- a/vortex-file/src/segments/mod.rs +++ b/vortex-file/src/segments/mod.rs @@ -2,8 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod cache; +mod padding; mod source; pub(crate) mod writer; pub use cache::*; +pub use padding::*; pub use source::*; diff --git a/vortex-file/src/segments/padding/mod.rs b/vortex-file/src/segments/padding/mod.rs new file mode 100644 index 00000000000..2b0bc65eb7c --- /dev/null +++ b/vortex-file/src/segments/padding/mod.rs @@ -0,0 +1,184 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::LazyLock; + +use vortex_buffer::Alignment; + +/// The block size assumed by [`SegmentPadding::block_aligned`], matching the direct-I/O and page +/// granularity of every mainstream filesystem we target. +pub const DEFAULT_BLOCK_SIZE: Alignment = Alignment::new(4096); + +/// The default overhead budget for [`SegmentPadding::proportional`], bounding the padding the +/// writer may add to at most `1/64` (1.6%) of the segment bytes written. +pub const DEFAULT_MAX_OVERHEAD_RATIO: u32 = 64; + +/// Environment variable selecting the default [`SegmentPadding`] policy. +pub const SEGMENT_PADDING_ENV_VAR: &str = "VORTEX_SEGMENT_PADDING"; + +static PADDING_FROM_ENV: LazyLock = LazyLock::new(|| { + let Ok(value) = std::env::var(SEGMENT_PADDING_ENV_VAR) else { + return SegmentPadding::None; + }; + // An empty value reads as unset, so a workflow can template the variable in unconditionally. + if value.trim().is_empty() { + return SegmentPadding::None; + } + SegmentPadding::parse(&value).unwrap_or_else(|| { + tracing::warn!( + "ignoring unrecognised {SEGMENT_PADDING_ENV_VAR}={value}, \ + expected none, always, grouped, or proportional[:ratio]" + ); + SegmentPadding::None + }) +}); + +/// How the file writer positions segments relative to storage block boundaries. +/// +/// A segment is always padded enough to satisfy its own memory alignment, so that a reader can +/// hand the bytes straight to a typed array without copying. This policy controls whether the +/// writer pads *further*, up to a storage block boundary. +/// +/// ## Why this is a trade-off +/// +/// Block alignment does not make direct I/O possible — a reader can always widen an unaligned +/// request out to the enclosing blocks and slice the result — it only removes the widening. That +/// saves at most one block of over-read per *physical* read, while padding costs up to one block +/// per *segment*. Because reads are coalesced across many segments, aligning every segment +/// usually costs far more bytes than it saves. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum SegmentPadding { + /// Pack segments contiguously, padding only as far as each segment's own alignment requires. + #[default] + None, + /// Start every segment on a `block` boundary. + /// + /// Costs on average half a block per segment, regardless of how small the segment is. + Always { + /// The storage block size to align to. + block: Alignment, + }, + /// Keep every sub-block segment inside a single `block`, packing runs of them together. + /// + /// A segment smaller than `block` stays where it is when it fits in what is left of the + /// current block, and moves to the next block only when it would straddle one. Runs of small + /// consecutive segments therefore share a block instead of each burning a whole one, and each + /// still reads in the single block [`Self::Always`] would have given it. + /// + /// A segment at least a block long is never moved. It already spans several blocks, so the + /// one extra block a straddle can cost it is a small fraction of its read, and far cheaper + /// than the up-to-a-block of padding that aligning it would add to the file. + Grouped { + /// The storage block size to align to. + block: Alignment, + }, + /// Start a segment on a `block` boundary only when the padding is worth it. + /// + /// A segment is aligned when its padding is at most `1/max_overhead_ratio` of the segment's + /// own length, which bounds the padding added across the whole file to `1/max_overhead_ratio` + /// of the segment bytes written. Large segments — the ones that dominate bytes read — are + /// nearly always aligned, while small segments are packed contiguously as before. + Proportional { + /// The storage block size to align to. + block: Alignment, + /// The reciprocal of the padding budget. Larger values pad less. + max_overhead_ratio: u32, + }, +} + +impl SegmentPadding { + /// Start every segment on a [`DEFAULT_BLOCK_SIZE`] boundary. + pub const fn block_aligned() -> Self { + Self::Always { + block: DEFAULT_BLOCK_SIZE, + } + } + + /// Keep every segment below [`DEFAULT_BLOCK_SIZE`] within a single block of that size. + pub const fn grouped() -> Self { + Self::Grouped { + block: DEFAULT_BLOCK_SIZE, + } + } + + /// Block-align segments within the default [`DEFAULT_MAX_OVERHEAD_RATIO`] padding budget. + pub const fn proportional() -> Self { + Self::Proportional { + block: DEFAULT_BLOCK_SIZE, + max_overhead_ratio: DEFAULT_MAX_OVERHEAD_RATIO, + } + } + + /// The policy named by `VORTEX_SEGMENT_PADDING`, or [`Self::None`] if it is unset. + /// + /// Accepts `none`, `always`, `grouped`, and `proportional[:ratio]`, all at the + /// [`DEFAULT_BLOCK_SIZE`] block size, so a deployment or benchmark can switch policies + /// without recompiling. Other block sizes are reachable only through the variants. + pub fn from_env() -> Self { + *PADDING_FROM_ENV + } + + fn parse(value: &str) -> Option { + let value = value.trim(); + let (name, ratio) = match value.split_once(':') { + Some((name, ratio)) => (name, Some(ratio.parse().ok()?)), + None => (value, None), + }; + match (name, ratio) { + ("none", None) => Some(Self::None), + ("always", None) => Some(Self::block_aligned()), + ("grouped", None) => Some(Self::grouped()), + ("proportional", None) => Some(Self::proportional()), + // A zero budget would pad everything, which `always` already says more clearly. + ("proportional", Some(ratio)) if ratio > 0 => Some(Self::Proportional { + block: DEFAULT_BLOCK_SIZE, + max_overhead_ratio: ratio, + }), + _ => None, + } + } + + /// The number of padding bytes to insert before a segment of `length` bytes requiring + /// `alignment`, when the writer is positioned at `byte_offset`. + pub fn padding(self, byte_offset: u64, length: u64, alignment: Alignment) -> u64 { + let required = pad_to(byte_offset, alignment); + match self { + Self::None => required, + // A segment whose own alignment exceeds the block size still has to satisfy it, so + // align to whichever is larger. Both are powers of two, so the larger subsumes both. + Self::Always { block } => pad_to(byte_offset, block.max(alignment)), + // Only a sub-block segment is worth moving: straddling a boundary doubles the blocks + // it touches, from one to two. A segment at least a block long already spans several, + // so the extra block a straddle costs it is a small fraction, not worth a whole block + // of padding. Within that, a segment that fits in the current block's remainder rides + // along for free. + Self::Grouped { block } => { + let block_size = block.as_usize() as u64; + let offset_in_block = (byte_offset + required) % block_size; + if length >= block_size || offset_in_block + length <= block_size { + required + } else { + pad_to(byte_offset, block.max(alignment)) + } + } + Self::Proportional { + block, + max_overhead_ratio, + } => { + let aligned = pad_to(byte_offset, block.max(alignment)); + if aligned.saturating_mul(u64::from(max_overhead_ratio)) <= length { + aligned + } else { + required + } + } + } + } +} + +fn pad_to(byte_offset: u64, alignment: Alignment) -> u64 { + byte_offset.next_multiple_of(alignment.as_usize() as u64) - byte_offset +} + +#[cfg(test)] +mod tests; diff --git a/vortex-file/src/segments/padding/tests.rs b/vortex-file/src/segments/padding/tests.rs new file mode 100644 index 00000000000..edea2cf8029 --- /dev/null +++ b/vortex-file/src/segments/padding/tests.rs @@ -0,0 +1,392 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; + +use super::*; + +const BLOCK: Alignment = Alignment::new(4096); + +#[rstest] +// Contiguous packing pads only up to the segment's own alignment. +#[case(SegmentPadding::None, 10, 1 << 20, Alignment::new(8), 6)] +#[case(SegmentPadding::None, 16, 1 << 20, Alignment::new(8), 0)] +// Always pads to the block, however small the segment. +#[case(SegmentPadding::block_aligned(), 10, 4, Alignment::new(8), 4086)] +#[case(SegmentPadding::block_aligned(), 4096, 4, Alignment::new(8), 0)] +// A segment demanding more than a block is aligned to its own requirement. +#[case(SegmentPadding::block_aligned(), 10, 4, Alignment::new(1 << 16), 65526)] +// Grouping leaves a segment that fits in the rest of the current block where it is. +#[case(SegmentPadding::grouped(), 10, 4, Alignment::new(8), 6)] +#[case(SegmentPadding::grouped(), 4000, 96, Alignment::new(8), 0)] +// One byte more than the block has room for, and it moves to the next block. +#[case(SegmentPadding::grouped(), 4000, 97, Alignment::new(8), 96)] +// Anything at least a block long is left packed: a straddle costs it far less than padding. +#[case(SegmentPadding::grouped(), 10, 1 << 20, Alignment::new(8), 6)] +#[case(SegmentPadding::grouped(), 4000, 4096, Alignment::new(8), 0)] +// A segment demanding more than a block still only pays its own alignment. +#[case(SegmentPadding::grouped(), 10, 4, Alignment::new(1 << 16), 65526)] +// Proportional pads a large segment, but leaves a small one packed. +#[case(SegmentPadding::proportional(), 10, 1 << 20, Alignment::new(8), 4086)] +#[case(SegmentPadding::proportional(), 10, 4, Alignment::new(8), 6)] +// Exactly on budget: 4086 * 64 <= 261_504. +#[case(SegmentPadding::proportional(), 10, 261_504, Alignment::new(8), 4086)] +#[case(SegmentPadding::proportional(), 10, 261_503, Alignment::new(8), 6)] +fn pads_segments( + #[case] padding: SegmentPadding, + #[case] byte_offset: u64, + #[case] length: u64, + #[case] alignment: Alignment, + #[case] expected: u64, +) { + assert_eq!(padding.padding(byte_offset, length, alignment), expected); +} + +#[rstest] +#[case(SegmentPadding::block_aligned())] +#[case(SegmentPadding::grouped())] +#[case(SegmentPadding::proportional())] +#[case(SegmentPadding::None)] +fn padding_always_satisfies_the_segments_own_alignment(#[case] padding: SegmentPadding) { + for byte_offset in 0..300u64 { + for alignment in [1, 2, 8, 64, 256] { + let alignment = Alignment::new(alignment); + let offset = byte_offset + padding.padding(byte_offset, 1 << 20, alignment); + assert!(alignment.is_offset_aligned(offset as usize)); + } + } +} + +/// The proportional budget must hold across a whole file, not just per segment. +#[test] +fn proportional_padding_is_bounded_by_the_overhead_budget() { + let padding = SegmentPadding::proportional(); + let mut offset = 0u64; + let mut padded = 0u64; + let mut data = 0u64; + for i in 0..10_000u64 { + // A spread of segment sizes from a few bytes to a few MiB. + let length = 3 + (i * 7919) % (4 << 20); + let pad = padding.padding(offset, length, Alignment::new(8)); + offset += pad + length; + padded += pad; + data += length; + } + assert!( + padded * u64::from(DEFAULT_MAX_OVERHEAD_RATIO) <= data, + "padding {padded} exceeded the 1/{DEFAULT_MAX_OVERHEAD_RATIO} budget of {data} bytes" + ); +} + +/// Block alignment must never cost more than a block per segment. +/// The whole point of grouping: sub-block segments touch one block each, as they would under +/// block alignment, for a fraction of the padding. +#[test] +fn grouping_reads_like_block_alignment_for_less_padding() { + let (mut grouped_offset, mut always_offset) = (0u64, 0u64); + let (mut grouped_pad, mut always_pad) = (0u64, 0u64); + for i in 0..10_000u64 { + // Sub-block segments, which is where the two policies are supposed to agree on reads. + let length = 1 + (i * 7919) % ((1 << 12) - 1); + let alignment = Alignment::new(1 << (i % 5)); + + let pad = SegmentPadding::grouped().padding(grouped_offset, length, alignment); + grouped_offset += pad; + grouped_pad += pad; + let pad = SegmentPadding::block_aligned().padding(always_offset, length, alignment); + always_offset += pad; + always_pad += pad; + + assert_eq!( + blocks_spanned(grouped_offset, length), + blocks_spanned(always_offset, length), + "segment {i} of {length} bytes spans a different number of blocks" + ); + + grouped_offset += length; + always_offset += length; + } + // Grouping pads a fraction of what block-aligning does on this distribution. + assert!( + grouped_pad * 2 < always_pad, + "grouping padded {grouped_pad} against {always_pad} block-aligned" + ); +} + +/// A grouped segment never straddles a block it could have fit inside. +#[test] +fn grouping_never_splits_a_segment_that_fits_in_a_block() { + let mut offset = 0u64; + for i in 0..10_000u64 { + let length = 1 + (i * 4099) % (1 << 13); + offset += SegmentPadding::grouped().padding(offset, length, Alignment::new(8)); + if length < BLOCK.as_usize() as u64 { + assert_eq!(blocks_spanned(offset, length), 1, "segment {i} was split"); + } + offset += length; + } +} + +fn blocks_spanned(offset: u64, length: u64) -> u64 { + let block = BLOCK.as_usize() as u64; + (offset % block + length).div_ceil(block) +} + +#[rstest] +#[case("none", SegmentPadding::None)] +#[case("always", SegmentPadding::block_aligned())] +#[case("grouped", SegmentPadding::grouped())] +#[case("proportional", SegmentPadding::proportional())] +#[case(" grouped ", SegmentPadding::grouped())] +#[case( + "proportional:8", + SegmentPadding::Proportional { block: BLOCK, max_overhead_ratio: 8 } +)] +fn parses_a_padding_policy(#[case] value: &str, #[case] expected: SegmentPadding) { + assert_eq!(SegmentPadding::parse(value), Some(expected)); +} + +#[rstest] +#[case("")] +#[case("grouped:4096")] +#[case("aligned")] +#[case("proportional:0")] +#[case("proportional:-1")] +#[case("proportional:")] +fn rejects_an_unknown_padding_policy(#[case] value: &str) { + assert_eq!(SegmentPadding::parse(value), None); +} + +#[test] +fn block_alignment_costs_at_most_one_block_per_segment() { + let padding = SegmentPadding::block_aligned(); + for byte_offset in [0u64, 1, 7, 4095, 4096, 1 << 20, (1 << 20) + 3] { + let pad = padding.padding(byte_offset, 1024, Alignment::new(8)); + assert!(pad < BLOCK.as_usize() as u64); + assert_eq!((byte_offset + pad) % BLOCK.as_usize() as u64, 0); + } +} + +mod end_to_end { + use std::sync::Arc; + use std::sync::LazyLock; + + use tempfile::TempDir; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::assert_arrays_eq; + use vortex_array::memory::BufferAllocatorRef; + use vortex_array::stream::ArrayStreamExt; + use vortex_buffer::Alignment; + use vortex_error::VortexResult; + use vortex_io::runtime::Handle; + use vortex_io::session::RuntimeSession; + use vortex_io::session::RuntimeSessionExt; + use vortex_io::std_file::FileReadAt; + use vortex_io::std_file::FileReadAtOptions; + use vortex_io::std_file::FileWrite; + use vortex_layout::session::LayoutSession; + use vortex_session::VortexSession; + + use super::*; + use crate::OpenOptionsSessionExt; + use crate::WriteOptionsSessionExt; + use crate::footer::SegmentSpec; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = array_session() + .with::() + .with::(); + crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); + session + }); + + /// Enough chunks and columns to produce a spread of small and large segments. + fn sample() -> VortexResult { + let numbers = + ChunkedArray::from_iter((0..8).map(|c| { + PrimitiveArray::from_iter((0..2000i64).map(|i| i * (c + 1))).into_array() + })) + .into_array(); + let strings = ChunkedArray::from_iter((0..8).map(|c| { + VarBinArray::from_iter( + (0..2000).map(|i| Some(format!("chunk-{c}-row-{i}"))), + vortex_array::dtype::DType::Utf8(vortex_array::dtype::Nullability::Nullable), + ) + .into_array() + })) + .into_array(); + Ok(StructArray::from_fields(&[("numbers", numbers), ("strings", strings)])?.into_array()) + } + + async fn write_to(path: &std::path::Path, padding: SegmentPadding) -> VortexResult<()> { + let write = FileWrite::create(path, SESSION.handle()).await?; + SESSION + .write_options() + .with_segment_padding(padding) + .write(write, sample()?.to_array_stream()) + .await?; + Ok(()) + } + + async fn segment_specs(path: &std::path::Path) -> VortexResult> { + let file = SESSION.open_options().open_path(path).await?; + Ok(Arc::clone(file.footer().segment_map())) + } + + fn block_of(spec: &SegmentSpec) -> u64 { + DEFAULT_BLOCK_SIZE.max(spec.alignment).as_usize() as u64 + } + + #[tokio::test] + async fn block_alignment_aligns_every_segment() -> VortexResult<()> { + let dir = TempDir::new()?; + let path = dir.path().join("aligned.vortex"); + write_to(&path, SegmentPadding::block_aligned()).await?; + + let specs = segment_specs(&path).await?; + assert!(!specs.is_empty()); + for spec in specs.iter() { + assert_eq!( + spec.offset % block_of(spec), + 0, + "segment at {} is not block aligned", + spec.offset + ); + } + Ok(()) + } + + #[tokio::test] + async fn contiguous_packing_leaves_segments_unaligned() -> VortexResult<()> { + let dir = TempDir::new()?; + let path = dir.path().join("packed.vortex"); + write_to(&path, SegmentPadding::None).await?; + + let specs = segment_specs(&path).await?; + assert!( + specs.iter().any(|spec| spec.offset % block_of(spec) != 0), + "expected contiguously packed segments to straddle block boundaries" + ); + Ok(()) + } + + /// Proportional padding must align the segments that dominate the bytes read while leaving the + /// long tail of small segments packed. + #[tokio::test] + async fn proportional_padding_aligns_only_the_large_segments() -> VortexResult<()> { + let dir = TempDir::new()?; + let path = dir.path().join("proportional.vortex"); + write_to(&path, SegmentPadding::proportional()).await?; + + let specs = segment_specs(&path).await?; + let budget = u64::from(DEFAULT_MAX_OVERHEAD_RATIO); + for spec in specs.iter() { + let block = block_of(spec); + if spec.offset % block != 0 { + // Only a segment too small to afford a block boundary may be left unaligned. + let required = (block - spec.offset % block) % block; + assert!( + required * budget > u64::from(spec.length), + "segment of {} bytes at {} could have afforded {required} bytes of padding", + spec.length, + spec.offset + ); + } + } + Ok(()) + } + + async fn read_back( + path: &std::path::Path, + handle: Handle, + options: FileReadAtOptions, + ) -> VortexResult { + let reader = FileReadAt::open_with_options( + path, + handle, + BufferAllocatorRef::statically_allocated(), + options, + )?; + SESSION + .open_options() + .open(Arc::new(reader)) + .await? + .scan()? + .into_array_stream()? + .read_all() + .await + } + + /// Every padding policy must round-trip identically under both buffered and direct reads, + /// including the contiguously packed layout that files written before this option used. + #[rstest] + #[case(SegmentPadding::None)] + #[case(SegmentPadding::block_aligned())] + #[case(SegmentPadding::grouped())] + #[case(SegmentPadding::proportional())] + #[tokio::test] + async fn direct_reads_round_trip_every_padding_policy( + #[case] padding: SegmentPadding, + ) -> VortexResult<()> { + let dir = TempDir::new()?; + let path = dir.path().join("data.vortex"); + write_to(&path, padding).await?; + + let handle = SESSION.handle(); + let direct = { + #[cfg(target_os = "linux")] + { + FileReadAtOptions::default().with_direct_io() + } + #[cfg(not(target_os = "linux"))] + { + FileReadAtOptions::default() + } + }; + + let mut ctx = SESSION.create_execution_ctx(); + let expected = read_back(&path, handle.clone(), FileReadAtOptions::default()).await?; + let actual = read_back(&path, handle, direct).await?; + assert_arrays_eq!(expected, actual, &mut ctx); + Ok(()) + } + + /// Block alignment must be a pure layout change: the same bytes, at padded offsets. + #[tokio::test] + async fn block_alignment_only_moves_segments() -> VortexResult<()> { + let dir = TempDir::new()?; + let packed = dir.path().join("packed.vortex"); + let aligned = dir.path().join("aligned.vortex"); + write_to(&packed, SegmentPadding::None).await?; + write_to(&aligned, SegmentPadding::block_aligned()).await?; + + let packed_specs = segment_specs(&packed).await?; + let aligned_specs = segment_specs(&aligned).await?; + assert_eq!(packed_specs.len(), aligned_specs.len()); + for (packed, aligned) in packed_specs.iter().zip(aligned_specs.iter()) { + assert_eq!(packed.length, aligned.length); + assert_eq!(packed.alignment, aligned.alignment); + } + + assert!( + std::fs::metadata(&aligned)?.len() > std::fs::metadata(&packed)?.len(), + "block alignment should cost padding bytes" + ); + Ok(()) + } + + #[test] + fn write_options_default_to_contiguous_packing() { + assert_eq!(SegmentPadding::default(), SegmentPadding::None); + } + + const _: () = assert!(DEFAULT_BLOCK_SIZE.as_usize() == 4096); + const _: Alignment = DEFAULT_BLOCK_SIZE; +} diff --git a/vortex-file/src/segments/writer.rs b/vortex-file/src/segments/writer.rs index 1673ce40101..42b388da65d 100644 --- a/vortex-file/src/segments/writer.rs +++ b/vortex-file/src/segments/writer.rs @@ -16,19 +16,26 @@ use vortex_layout::segments::SegmentSink; use vortex_layout::sequence::SequenceId; use crate::footer::SegmentSpec; +use crate::segments::SegmentPadding; pub struct BufferedSegmentSink { buffers: kanal::AsyncSender, byte_offset: AtomicU64, segment_specs: Mutex>, + padding: SegmentPadding, } impl BufferedSegmentSink { - pub fn new(send: kanal::AsyncSender, byte_offset: u64) -> Self { + pub fn new( + send: kanal::AsyncSender, + byte_offset: u64, + padding: SegmentPadding, + ) -> Self { Self { buffers: send, byte_offset: AtomicU64::new(byte_offset), segment_specs: Default::default(), + padding, } } @@ -70,7 +77,9 @@ impl SegmentSink for BufferedSegmentSink { // Add any padding required to align the segment. let byte_offset = self.byte_offset.load(Ordering::Relaxed); - let padding = byte_offset.next_multiple_of(alignment.as_usize() as u64) - byte_offset; + let padding = self + .padding + .padding(byte_offset, u64::from(length), alignment); let offset = byte_offset + padding; specs.push(SegmentSpec { offset, diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index ec45653f5c1..607f75cb2c1 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -66,6 +66,7 @@ use crate::counting::CountingVortexWrite; use crate::footer::FileStatistics; use crate::footer::MAX_METADATA_KEY_BYTES; use crate::footer::MAX_METADATA_SEGMENTS; +use crate::segments::SegmentPadding; use crate::segments::writer::BufferedSegmentSink; /// Configure a new writer, which can eventually be used to write an [`ArrayStream`] into a sink @@ -86,6 +87,7 @@ pub struct VortexWriteOptions { max_variable_length_statistics_size: usize, file_statistics: Vec, metadata: HashMap, + segment_padding: SegmentPadding, } /// Extension trait for constructing [`VortexWriteOptions`] from a session. @@ -109,9 +111,21 @@ impl VortexWriteOptions { file_statistics: PRUNING_STATS.to_vec(), max_variable_length_statistics_size: 64, metadata: HashMap::default(), + segment_padding: SegmentPadding::from_env(), } } + /// Control how segments are positioned relative to storage block boundaries. + /// + /// By default segments are packed contiguously, padded only as far as their own memory + /// alignment requires, unless `VORTEX_SEGMENT_PADDING` names another policy. Block-aligning + /// them lets a direct-I/O reader issue each read without widening it to the enclosing blocks, + /// at the cost of the padding bytes. See [`SegmentPadding`] for the trade-off. + pub fn with_segment_padding(mut self, segment_padding: SegmentPadding) -> Self { + self.segment_padding = segment_padding; + self + } + /// Replace the default layout strategy with the provided one. /// /// The strategy controls repartitioning, statistics layout, compression, and leaf segment @@ -285,7 +299,11 @@ impl VortexWriteOptions { // Create a channel to send buffers from the segment sink to the output stream. let (send, recv) = kanal::bounded_async(1); - let segments = Arc::new(BufferedSegmentSink::new(send, position)); + let segments = Arc::new(BufferedSegmentSink::new( + send, + position, + self.segment_padding, + )); // We spawn the layout future so it is driven in the background while we write the // buffer stream, so we don't need to poll it until all buffers have been drained. diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index 3905084d85a..3794eeddc91 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -45,6 +45,9 @@ vortex-utils = { workspace = true } [target.'cfg(unix)'.dependencies] custom-labels = { workspace = true } +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true } + [target.'cfg(not(target_arch = "wasm32"))'.dependencies] # Smol is our default impl, so we don't want it to be optional, but it cannot be part of wasm smol = { workspace = true } diff --git a/vortex-io/src/std_file/direct.rs b/vortex-io/src/std_file/direct.rs new file mode 100644 index 00000000000..2b44770f8f7 --- /dev/null +++ b/vortex-io/src/std_file/direct.rs @@ -0,0 +1,285 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Direct (page-cache bypassing) reads of local files. +//! +//! Linux imposes three constraints on `O_DIRECT` reads: the file offset, the transfer length, and +//! the address of the user-space buffer must each be aligned. The required alignments are reported +//! per-file by `statx(STATX_DIOALIGN)`; when the kernel or filesystem does not report them we fall +//! back to a page-sized alignment. +//! +//! Vortex segments are aligned to their element width, not to a block boundary, so a logical read +//! almost never satisfies these constraints on its own. [`DirectIoConstraints::widen`] grows the +//! request out to the enclosing block boundaries and records where the requested bytes sit inside +//! the widened window, so the caller can slice them back out after the transfer. This is what lets +//! direct I/O read files written by any Vortex version, without a format change. + +use std::fs::File; +use std::io; +use std::ops::Range; +use std::os::unix::fs::FileExt; +use std::path::Path; + +use rustix::fs::AtFlags; +use rustix::fs::Mode; +use rustix::fs::OFlags; +use rustix::fs::StatxFlags; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +/// Conservative direct-I/O alignment used when Linux cannot report the filesystem constraints. +/// +/// A page-sized fallback is accepted by common block devices and filesystems. If the actual +/// requirement is stricter, the read fails with the underlying `EINVAL`. +pub const FALLBACK_DIRECT_IO_ALIGNMENT: usize = 4096; + +/// Open `path` read-only with the page cache bypassed. +pub fn open_direct(path: &Path) -> io::Result { + Ok(File::from( + rustix::fs::open( + path, + OFlags::RDONLY | OFlags::CLOEXEC | OFlags::DIRECT, + Mode::empty(), + ) + .map_err(io::Error::from)?, + )) +} + +/// The alignment a filesystem requires of direct-I/O reads against a particular file. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DirectIoConstraints { + memory_alignment: usize, + offset_alignment: usize, +} + +impl Default for DirectIoConstraints { + fn default() -> Self { + Self { + memory_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, + offset_alignment: FALLBACK_DIRECT_IO_ALIGNMENT, + } + } +} + +impl DirectIoConstraints { + /// Ask the kernel for the direct-I/O constraints that apply to `file`. + /// + /// Falls back to [`FALLBACK_DIRECT_IO_ALIGNMENT`] when `statx` is unavailable or does not + /// report `STATX_DIOALIGN` (Linux before 6.1, or a filesystem that does not implement it). + pub fn probe(file: &File) -> VortexResult { + let Ok(stat) = rustix::fs::statx( + file, + c"", + AtFlags::EMPTY_PATH | AtFlags::STATX_DONT_SYNC, + StatxFlags::DIOALIGN, + ) else { + return Ok(Self::default()); + }; + if stat.stx_mask & StatxFlags::DIOALIGN.bits() == 0 { + return Ok(Self::default()); + } + + let Ok(memory_alignment) = usize::try_from(stat.stx_dio_mem_align) else { + return Ok(Self::default()); + }; + let Ok(offset_alignment) = usize::try_from(stat.stx_dio_offset_align) else { + return Ok(Self::default()); + }; + if memory_alignment == 0 || offset_alignment == 0 { + return Ok(Self::default()); + } + vortex_ensure!( + memory_alignment.is_power_of_two(), + "direct I/O memory alignment must be a power of two, got {memory_alignment}" + ); + vortex_ensure!( + offset_alignment.is_power_of_two(), + "direct I/O offset alignment must be a power of two, got {offset_alignment}" + ); + + Ok(Self { + memory_alignment, + offset_alignment, + }) + } + + /// Required alignment of the address of the user-space I/O buffer. + pub fn memory_alignment(&self) -> usize { + self.memory_alignment + } + + /// Required alignment of both the file offset and the I/O length. + pub fn offset_alignment(&self) -> usize { + self.offset_alignment + } + + /// Widen `offset..offset + length` out to the enclosing direct-I/O block boundaries. + pub fn widen(&self, offset: u64, length: usize) -> VortexResult { + direct_io_range(offset, length, self.offset_alignment) + } + + /// Read at least `required_bytes` into `buffer`, returning the number of bytes initialized. + /// + /// `buffer` must be aligned to [`memory_alignment`][Self::memory_alignment] and its length + /// must be a multiple of [`offset_alignment`][Self::offset_alignment], as produced by + /// [`widen`][Self::widen]. Reads past the end of the file return short, which is expected for + /// the final block of a file whose length is not a multiple of the block size. + pub fn read_at( + &self, + file: &File, + buffer: &mut [u8], + offset: u64, + required_bytes: usize, + ) -> io::Result { + let mut initialized = 0; + while initialized < required_bytes { + let initialized_u64 = u64::try_from(initialized) + .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow"))?; + let read_offset = offset.checked_add(initialized_u64).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "read offset overflow") + })?; + let bytes_read = match file.read_at(&mut buffer[initialized..], 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 {initialized} bytes, but {required_bytes} bytes were required" + ), + )); + } + initialized += bytes_read; + // A resumed read must itself start on an aligned boundary, so a short read that does + // not land on one leaves us unable to issue the remainder. + if initialized < required_bytes + && (!initialized.is_multiple_of(self.offset_alignment) + || !initialized.is_multiple_of(self.memory_alignment)) + { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!( + "direct read returned an unaligned short read of {initialized} bytes before the required {required_bytes} bytes" + ), + )); + } + } + + Ok(initialized) + } +} + +/// A logical read widened to satisfy direct-I/O offset and length alignment. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DirectIoRange { + /// Block-aligned file offset to read from. + pub read_offset: u64, + /// Block-aligned number of bytes to read. + pub read_length: usize, + /// Position of the originally requested bytes within the widened read. + pub requested_range: Range, +} + +fn direct_io_range(offset: u64, length: usize, alignment: usize) -> VortexResult { + vortex_ensure!(alignment > 0, "direct I/O alignment must be non-zero"); + if length == 0 { + return Ok(DirectIoRange { + read_offset: offset, + read_length: 0, + requested_range: 0..0, + }); + } + + let alignment_u64 = u64::try_from(alignment)?; + let length_u64 = u64::try_from(length)?; + let requested_end = offset.checked_add(length_u64).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + let read_offset = offset - offset % alignment_u64; + let read_end = requested_end + .checked_next_multiple_of(alignment_u64) + .ok_or_else(|| vortex_err!("direct I/O aligned end overflow"))?; + let read_length = usize::try_from(read_end - read_offset)?; + let slice_start = usize::try_from(offset - read_offset)?; + let slice_end = slice_start.checked_add(length).ok_or_else(|| { + vortex_err!("direct I/O range overflow: offset={offset}, length={length}") + })?; + + Ok(DirectIoRange { + read_offset, + read_length, + requested_range: slice_start..slice_end, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case(0, 0, 4096, 0, 0, 0)] + #[case(5, 0, 4096, 5, 0, 0)] + #[case(5, 10, 4096, 0, 4096, 5)] + #[case(4090, 20, 4096, 0, 8192, 4090)] + #[case(4096, 4096, 4096, 4096, 4096, 0)] + #[case(513, 1, 512, 512, 512, 1)] + #[case(4096, 8193, 4096, 4096, 12288, 0)] + fn widens_direct_read_to_block_boundaries( + #[case] offset: u64, + #[case] length: usize, + #[case] alignment: usize, + #[case] expected_offset: u64, + #[case] expected_length: usize, + #[case] expected_prefix: usize, + ) -> VortexResult<()> { + assert_eq!( + direct_io_range(offset, length, alignment)?, + DirectIoRange { + read_offset: expected_offset, + read_length: expected_length, + requested_range: expected_prefix..expected_prefix + length, + } + ); + Ok(()) + } + + #[rstest] + #[case(u64::MAX, 2, 4096)] + #[case(0, 1, 0)] + fn rejects_invalid_direct_read_range( + #[case] offset: u64, + #[case] length: usize, + #[case] alignment: usize, + ) { + assert!(direct_io_range(offset, length, alignment).is_err()); + } + + #[test] + fn aligned_ranges_cover_requested_bytes() -> VortexResult<()> { + for alignment in [512, 4096] { + for offset in 0..alignment * 2 { + for length in [0, 1, alignment - 1, alignment, alignment + 1] { + let range = direct_io_range(offset as u64, length, alignment)?; + if length == 0 { + assert_eq!(range.read_length, 0); + continue; + } + + assert_eq!(range.read_offset % alignment as u64, 0); + assert_eq!(range.read_length % alignment, 0); + assert_eq!(range.requested_range.len(), length); + assert!(range.requested_range.end <= range.read_length); + assert_eq!( + range.read_offset + range.requested_range.start as u64, + offset as u64 + ); + } + } + } + Ok(()) + } +} diff --git a/vortex-io/src/std_file/mod.rs b/vortex-io/src/std_file/mod.rs index 3345089a023..5f051fcf267 100644 --- a/vortex-io/src/std_file/mod.rs +++ b/vortex-io/src/std_file/mod.rs @@ -1,10 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +#[cfg(target_os = "linux")] +mod direct; mod filesystem; mod read_at; mod write; +#[cfg(target_os = "linux")] +pub use direct::*; pub use filesystem::*; pub use read_at::*; pub use write::*; diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs deleted file mode 100644 index 21aac88923f..00000000000 --- a/vortex-io/src/std_file/read_at.rs +++ /dev/null @@ -1,139 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fs::File; -use std::io; -#[cfg(all(not(unix), not(windows)))] -use std::io::Read; -#[cfg(all(not(unix), not(windows)))] -use std::io::Seek; -#[cfg(unix)] -use std::os::unix::fs::FileExt; -#[cfg(windows)] -use std::os::windows::fs::FileExt; -use std::path::Path; -use std::sync::Arc; - -use futures::FutureExt; -use futures::future::BoxFuture; -use vortex_array::buffer::BufferHandle; -use vortex_array::memory::BufferAllocatorRef; -use vortex_buffer::Alignment; -use vortex_error::VortexResult; - -use crate::CoalesceConfig; -use crate::VortexReadAt; -use crate::runtime::Handle; - -/// Read exactly `buffer.len()` bytes from `file` starting at `offset`. -/// This is a platform-specific helper that uses the most efficient method available. -#[cfg(not(target_arch = "wasm32"))] -pub fn read_exact_at(file: &File, buffer: &mut [u8], offset: u64) -> io::Result<()> { - #[cfg(unix)] - { - file.read_exact_at(buffer, offset) - } - #[cfg(windows)] - { - let mut bytes_read = 0; - while bytes_read < buffer.len() { - let read = file.seek_read(&mut buffer[bytes_read..], offset + bytes_read as u64)?; - if read == 0 { - return Err(io::Error::new( - io::ErrorKind::UnexpectedEof, - "failed to fill whole buffer", - )); - } - bytes_read += read; - } - Ok(()) - } - #[cfg(all(not(unix), not(windows)))] - { - use std::io::SeekFrom; - let mut file_ref = file; - file_ref.seek(SeekFrom::Start(offset))?; - file_ref.read_exact(buffer) - } -} - -/// Default number of concurrent requests to allow for local file I/O. -pub const DEFAULT_CONCURRENCY: usize = 32; - -/// An adapter type wrapping a [`File`] to implement [`VortexReadAt`]. -pub struct FileReadAt { - uri: Arc, - file: Arc, - handle: Handle, - allocator: BufferAllocatorRef, -} - -impl FileReadAt { - /// Open a file for reading. - pub fn open(path: impl AsRef, handle: Handle) -> VortexResult { - Self::open_with_allocator(path, handle, BufferAllocatorRef::statically_allocated()) - } - - /// Open a file for reading using a custom writable buffer allocator. - pub fn open_with_allocator( - path: impl AsRef, - handle: Handle, - allocator: BufferAllocatorRef, - ) -> VortexResult { - let path = path.as_ref(); - let uri = path.to_string_lossy().to_string().into(); - let file = Arc::new(File::open(path)?); - Ok(Self { - uri, - file, - handle, - allocator, - }) - } -} - -impl VortexReadAt for FileReadAt { - fn uri(&self) -> Option<&Arc> { - Some(&self.uri) - } - - fn coalesce_config(&self) -> Option { - Some(CoalesceConfig::file()) - } - - fn concurrency(&self) -> usize { - DEFAULT_CONCURRENCY - } - - fn size(&self) -> BoxFuture<'static, VortexResult> { - let file = Arc::clone(&self.file); - async move { - let metadata = file.metadata()?; - Ok(metadata.len()) - } - .boxed() - } - - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - let file = Arc::clone(&self.file); - let handle = self.handle.clone(); - let allocator = self.allocator.clone(); - async move { - handle - .spawn_blocking(move || { - let mut buffer = allocator.with_capacity_aligned::(length, alignment); - // SAFETY: read_exact_at initializes every byte before the buffer is frozen. - unsafe { buffer.set_len(length) }; - read_exact_at(&file, buffer.as_mut_slice(), offset)?; - Ok(BufferHandle::new_host(buffer.freeze())) - }) - .await - } - .boxed() - } -} diff --git a/vortex-io/src/std_file/read_at/mod.rs b/vortex-io/src/std_file/read_at/mod.rs new file mode 100644 index 00000000000..c24f41dec02 --- /dev/null +++ b/vortex-io/src/std_file/read_at/mod.rs @@ -0,0 +1,303 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fs::File; +use std::io; +#[cfg(all(not(unix), not(windows)))] +use std::io::Read; +#[cfg(all(not(unix), not(windows)))] +use std::io::Seek; +#[cfg(unix)] +use std::os::unix::fs::FileExt; +#[cfg(windows)] +use std::os::windows::fs::FileExt; +use std::path::Path; +use std::sync::Arc; +use std::sync::LazyLock; + +use futures::FutureExt; +use futures::future::BoxFuture; +use vortex_array::buffer::BufferHandle; +use vortex_array::memory::BufferAllocatorRef; +use vortex_buffer::Alignment; +use vortex_buffer::ByteBufferMut; +use vortex_error::VortexResult; + +use crate::CoalesceConfig; +use crate::VortexReadAt; +use crate::runtime::Handle; +#[cfg(target_os = "linux")] +use crate::std_file::direct::DirectIoConstraints; +#[cfg(target_os = "linux")] +use crate::std_file::direct::open_direct; + +#[cfg(test)] +mod tests; + +/// Read exactly `buffer.len()` bytes from `file` starting at `offset`. +/// This is a platform-specific helper that uses the most efficient method available. +#[cfg(not(target_arch = "wasm32"))] +pub fn read_exact_at(file: &File, buffer: &mut [u8], offset: u64) -> io::Result<()> { + #[cfg(unix)] + { + file.read_exact_at(buffer, offset) + } + #[cfg(windows)] + { + let mut bytes_read = 0; + while bytes_read < buffer.len() { + let read = file.seek_read(&mut buffer[bytes_read..], offset + bytes_read as u64)?; + if read == 0 { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "failed to fill whole buffer", + )); + } + bytes_read += read; + } + Ok(()) + } + #[cfg(all(not(unix), not(windows)))] + { + use std::io::SeekFrom; + let mut file_ref = file; + file_ref.seek(SeekFrom::Start(offset))?; + file_ref.read_exact(buffer) + } +} + +/// Default number of concurrent requests to allow for local file I/O. +pub const DEFAULT_CONCURRENCY: usize = 32; + +/// Environment variable that opts local file reads into direct I/O. +pub const DIRECT_IO_ENV_VAR: &str = "VORTEX_DIRECT_IO"; + +static DIRECT_IO_FROM_ENV: LazyLock = + LazyLock::new(|| std::env::var(DIRECT_IO_ENV_VAR).is_ok_and(|v| v == "1")); + +/// Options controlling how [`FileReadAt`] opens and reads a local file. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct FileReadAtOptions { + direct_io: bool, +} + +impl FileReadAtOptions { + /// Options taken from the process environment. + /// + /// Direct I/O is enabled when `VORTEX_DIRECT_IO=1`, letting a deployment or benchmark switch + /// between page-cached and direct reads without recompiling. + pub fn from_env() -> Self { + Self { + direct_io: *DIRECT_IO_FROM_ENV, + } + } + + /// Bypass the operating system page cache for file reads. + /// + /// This option is available only on Linux. Reads that do not meet the filesystem's direct-I/O + /// alignment requirements are widened to the enclosing block boundaries and sliced back to the + /// requested range, so files written by any Vortex version can be read this way. If the file + /// cannot be opened with `O_DIRECT`, the reader falls back to buffered reads. + #[cfg(target_os = "linux")] + pub fn with_direct_io(mut self) -> Self { + self.direct_io = true; + self + } + + /// Whether direct I/O was requested. + pub fn direct_io(&self) -> bool { + self.direct_io + } +} + +/// How [`FileReadAt`] transfers bytes out of the file. +enum ReadMode { + /// Ordinary `pread`, served through the operating system page cache. + Buffered, + /// `O_DIRECT` `pread`, bypassing the page cache. + #[cfg(target_os = "linux")] + Direct(DirectIoConstraints), +} + +/// An adapter type wrapping a [`File`] to implement [`VortexReadAt`]. +pub struct FileReadAt { + uri: Arc, + file: Arc, + mode: Arc, + handle: Handle, + allocator: BufferAllocatorRef, +} + +impl FileReadAt { + /// Open a file for reading. + pub fn open(path: impl AsRef, handle: Handle) -> VortexResult { + Self::open_with_allocator(path, handle, BufferAllocatorRef::statically_allocated()) + } + + /// Open a file for reading using a custom writable buffer allocator. + pub fn open_with_allocator( + path: impl AsRef, + handle: Handle, + allocator: BufferAllocatorRef, + ) -> VortexResult { + Self::open_with_options(path, handle, allocator, FileReadAtOptions::from_env()) + } + + /// Open a file for reading with explicit options. + pub fn open_with_options( + path: impl AsRef, + handle: Handle, + allocator: BufferAllocatorRef, + options: FileReadAtOptions, + ) -> VortexResult { + let path = path.as_ref(); + let uri: Arc = path.to_string_lossy().to_string().into(); + let (file, mode) = open_file(path, options, &uri)?; + Ok(Self { + uri, + file: Arc::new(file), + mode: Arc::new(mode), + handle, + allocator, + }) + } + + /// Whether this reader bypasses the operating system page cache. + pub fn is_direct(&self) -> bool { + match self.mode.as_ref() { + ReadMode::Buffered => false, + #[cfg(target_os = "linux")] + ReadMode::Direct(_) => true, + } + } +} + +#[cfg(target_os = "linux")] +fn open_file(path: &Path, options: FileReadAtOptions, uri: &str) -> io::Result<(File, ReadMode)> { + if !options.direct_io { + return Ok((File::open(path)?, ReadMode::Buffered)); + } + // Filesystems that cannot serve O_DIRECT reject the open outright, so degrade to buffered + // reads rather than failing to open a file we can read perfectly well. + match open_direct(path) { + Ok(file) => match DirectIoConstraints::probe(&file) { + Ok(constraints) => Ok((file, ReadMode::Direct(constraints))), + Err(err) => { + tracing::warn!( + "{uri}: direct I/O constraints unavailable, using buffered reads: {err}" + ); + Ok((File::open(path)?, ReadMode::Buffered)) + } + }, + Err(err) => { + tracing::warn!("{uri}: cannot open with O_DIRECT, using buffered reads: {err}"); + Ok((File::open(path)?, ReadMode::Buffered)) + } + } +} + +#[cfg(not(target_os = "linux"))] +fn open_file(path: &Path, options: FileReadAtOptions, uri: &str) -> io::Result<(File, ReadMode)> { + if options.direct_io { + tracing::warn!("{uri}: direct I/O is only supported on Linux, using buffered reads"); + } + Ok((File::open(path)?, ReadMode::Buffered)) +} + +/// Read `length` bytes at `offset`, returning a buffer aligned to `alignment`. +fn read_buffer( + file: &File, + mode: &ReadMode, + allocator: &BufferAllocatorRef, + offset: u64, + length: usize, + alignment: Alignment, +) -> VortexResult { + match mode { + ReadMode::Buffered => { + let mut buffer = allocator.with_capacity_aligned::(length, alignment); + // SAFETY: read_exact_at initializes every byte before the buffer is frozen. + unsafe { buffer.set_len(length) }; + read_exact_at(file, buffer.as_mut_slice(), offset)?; + Ok(BufferHandle::new_host(buffer.freeze())) + } + #[cfg(target_os = "linux")] + ReadMode::Direct(constraints) => { + if length == 0 { + let buffer = allocator.with_capacity_aligned::(0, alignment); + return Ok(BufferHandle::new_host(buffer.freeze())); + } + + let range = constraints.widen(offset, length)?; + // The pointer must satisfy the filesystem's requirement as well as the caller's, and + // over-aligning the base keeps the requested bytes aligned once sliced back out: a + // segment offset that is a multiple of `alignment` stays one relative to a block + // boundary, because both are powers of two and blocks are the larger of the pair. + let alloc_alignment = alignment.max(Alignment::new(constraints.memory_alignment())); + let mut buffer: ByteBufferMut = ByteBufferMut::with_capacity_aligned_in( + range.read_length, + alloc_alignment, + allocator.clone(), + ); + // SAFETY: the length is trimmed below to the prefix the kernel initialized. + unsafe { buffer.set_len(range.read_length) }; + let initialized = constraints.read_at( + file, + buffer.as_mut_slice(), + range.read_offset, + range.requested_range.end, + )?; + unsafe { buffer.set_len(initialized) }; + + Ok(BufferHandle::new_host( + buffer + .freeze() + .slice_unaligned(range.requested_range) + .aligned(alignment), + )) + } + } +} + +impl VortexReadAt for FileReadAt { + fn uri(&self) -> Option<&Arc> { + Some(&self.uri) + } + + fn coalesce_config(&self) -> Option { + Some(CoalesceConfig::file()) + } + + fn concurrency(&self) -> usize { + DEFAULT_CONCURRENCY + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + let file = Arc::clone(&self.file); + async move { + let metadata = file.metadata()?; + Ok(metadata.len()) + } + .boxed() + } + + fn read_at( + &self, + offset: u64, + length: usize, + alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + let file = Arc::clone(&self.file); + let mode = Arc::clone(&self.mode); + let handle = self.handle.clone(); + let allocator = self.allocator.clone(); + async move { + handle + .spawn_blocking(move || { + read_buffer(&file, &mode, &allocator, offset, length, alignment) + }) + .await + } + .boxed() + } +} diff --git a/vortex-io/src/std_file/read_at/tests.rs b/vortex-io/src/std_file/read_at/tests.rs new file mode 100644 index 00000000000..cb4a6efab70 --- /dev/null +++ b/vortex-io/src/std_file/read_at/tests.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg(feature = "tokio")] +#![expect(clippy::cast_possible_truncation, reason = "fixture sizes are small")] + +use std::io::Write; + +use rstest::rstest; +use tempfile::NamedTempFile; +use vortex_array::memory::BufferAllocatorRef; +use vortex_buffer::Alignment; +use vortex_error::VortexResult; + +use crate::VortexReadAt; +use crate::runtime::tokio::TokioRuntime; +use crate::std_file::FileReadAt; +use crate::std_file::FileReadAtOptions; + +/// A file length that is deliberately not a multiple of any plausible block size, so the final +/// block is partial and every direct read near the end returns short. +const FILE_LEN: usize = 40_000 + 37; + +fn contents() -> Vec { + (0..FILE_LEN).map(|i| (i % 251) as u8).collect() +} + +fn temp_file(contents: &[u8]) -> VortexResult { + let mut file = NamedTempFile::new()?; + file.write_all(contents)?; + file.flush()?; + Ok(file) +} + +fn direct_options() -> FileReadAtOptions { + #[cfg(target_os = "linux")] + { + FileReadAtOptions::default().with_direct_io() + } + #[cfg(not(target_os = "linux"))] + { + FileReadAtOptions::default() + } +} + +fn open(path: &std::path::Path, options: FileReadAtOptions) -> VortexResult { + FileReadAt::open_with_options( + path, + TokioRuntime::current(), + BufferAllocatorRef::statically_allocated(), + options, + ) +} + +#[test] +fn options_default_to_buffered_io() { + assert!(!FileReadAtOptions::default().direct_io()); +} + +#[cfg(target_os = "linux")] +#[test] +fn options_enable_direct_io() { + assert!(FileReadAtOptions::default().with_direct_io().direct_io()); +} + +/// Direct reads must return exactly the requested window, including when the request straddles or +/// sits inside a block, and when it runs up against a partial final block. +#[rstest] +#[case(0, 1)] +#[case(0, FILE_LEN)] +#[case(1, 4095)] +#[case(4095, 2)] +#[case(4096, 4096)] +#[case(511, 8193)] +#[case(FILE_LEN as u64 - 1, 1)] +#[case(FILE_LEN as u64 - 4097, 4097)] +#[case(35_000, FILE_LEN - 35_000)] +#[tokio::test] +async fn direct_reads_return_the_requested_bytes( + #[case] offset: u64, + #[case] length: usize, +) -> VortexResult<()> { + let expected = contents(); + let file = temp_file(&expected)?; + + let direct = open(file.path(), direct_options())?; + let buffered = open(file.path(), FileReadAtOptions::default())?; + + let direct = direct.read_at(offset, length, Alignment::none()).await?; + let buffered = buffered.read_at(offset, length, Alignment::none()).await?; + + let window = &expected[offset as usize..offset as usize + length]; + assert_eq!(direct.to_host().await.as_slice(), window); + assert_eq!(buffered.to_host().await.as_slice(), window); + Ok(()) +} + +/// Widening a read to block boundaries must not cost the caller their alignment: a segment stored +/// at a naturally aligned file offset is still naturally aligned once sliced out of the block. +#[rstest] +#[case(8, Alignment::new(8))] +#[case(4104, Alignment::new(8))] +#[case(256, Alignment::new(256))] +#[case(8192, Alignment::new(4096))] +#[tokio::test] +async fn direct_reads_preserve_requested_alignment( + #[case] offset: u64, + #[case] alignment: Alignment, +) -> VortexResult<()> { + let expected = contents(); + let file = temp_file(&expected)?; + let reader = open(file.path(), direct_options())?; + + let length = 1024; + let buffer = reader.read_at(offset, length, alignment).await?; + let host = buffer.to_host().await; + + assert_eq!(host.alignment(), alignment); + assert!(host.is_aligned(alignment)); + assert_eq!( + host.as_slice(), + &expected[offset as usize..offset as usize + length] + ); + Ok(()) +} + +#[tokio::test] +async fn direct_reads_past_the_end_of_the_file_fail() -> VortexResult<()> { + let file = temp_file(&contents())?; + let reader = open(file.path(), direct_options())?; + + assert!( + reader + .read_at(FILE_LEN as u64 - 8, 4096, Alignment::none()) + .await + .is_err() + ); + Ok(()) +} + +#[tokio::test] +async fn direct_reads_of_zero_length_are_empty() -> VortexResult<()> { + let file = temp_file(&contents())?; + let reader = open(file.path(), direct_options())?; + + let buffer = reader.read_at(37, 0, Alignment::new(8)).await?; + assert_eq!(buffer.len(), 0); + Ok(()) +} + +/// Direct I/O is unavailable on some filesystems (tmpfs, overlayfs) and on non-Linux platforms. +/// Opening must still succeed there, silently serving buffered reads. +#[tokio::test] +async fn opening_never_fails_when_direct_io_is_unavailable() -> VortexResult<()> { + let expected = contents(); + let file = temp_file(&expected)?; + let reader = open(file.path(), direct_options())?; + + let buffer = reader.read_at(0, 128, Alignment::none()).await?; + assert_eq!(buffer.to_host().await.as_slice(), &expected[..128]); + Ok(()) +} + +/// Wherever the platform can actually serve `O_DIRECT`, requesting it must take effect rather than +/// quietly degrading, otherwise the option would be untestable and unmeasurable. +#[cfg(target_os = "linux")] +#[tokio::test] +async fn direct_io_is_used_when_the_filesystem_supports_it() -> VortexResult<()> { + let file = temp_file(&contents())?; + if crate::std_file::open_direct(file.path()).is_err() { + return Ok(()); + } + + assert!(open(file.path(), direct_options())?.is_direct()); + assert!(!open(file.path(), FileReadAtOptions::default())?.is_direct()); + Ok(()) +}