Skip to content

Remove mutable state from Vortex open options #8879

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
onursatici merged 2 commits into from
Jul 21, 2026
Merged

Remove mutable state from Vortex open options #8879

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6877f1d
Remove mutable state from Vortex open options
onursatici Jul 21, 2026
c0086e5
remove manual impl Clone
onursatici Jul 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 35 additions & 43 deletions vortex-file/src/open.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
use std::sync::Arc;

use futures::executor::block_on;
use parking_lot::RwLock;
use vortex_array::dtype::DType;
use vortex_array::memory::MemorySessionExt;
use vortex_array::session::ArraySessionExt;
Expand Down Expand Up @@ -41,13 +40,19 @@ use crate::segments::RequestMetrics;

const INITIAL_READ_SIZE: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE;

struct FooterRead {
footer: Footer,
initial_segments: HashMap<SegmentId, ByteBuffer>,
}

/// Open options for a Vortex file reader.
///
/// Options are session-bound because opening a file may need array, layout, dtype, runtime, and
/// memory registries. Construct with [`OpenOptionsSessionExt::open_options`] for the common path.
///
/// The opener first resolves a [`Footer`], then creates a segment source for later scans. Known
/// file metadata can be supplied up front to avoid IO during footer discovery.
#[derive(Clone)]
pub struct VortexOpenOptions {
/// The session to use for opening the file.
session: VortexSession,
Expand All @@ -61,8 +66,6 @@ pub struct VortexOpenOptions {
dtype: Option<DType>,
/// An optional, externally provided, file layout.
footer: Option<Footer>,
/// The segments read during the initial read.
initial_read_segments: RwLock<HashMap<SegmentId, ByteBuffer>>,
/// A metrics registry for the file.
metrics_registry: Option<Arc<dyn MetricsRegistry>>,
/// Default labels applied to all the file's metrics
Expand All @@ -71,23 +74,6 @@ pub struct VortexOpenOptions {
cache_layout_reader: bool,
}

impl Clone for VortexOpenOptions {
fn clone(&self) -> Self {
Self {
session: self.session.clone(),
segment_cache: self.segment_cache.clone(),
initial_read_size: self.initial_read_size,
file_size: self.file_size,
dtype: self.dtype.clone(),
footer: self.footer.clone(),
initial_read_segments: RwLock::new(self.initial_read_segments.read().clone()),
metrics_registry: self.metrics_registry.clone(),
labels: self.labels.clone(),
cache_layout_reader: self.cache_layout_reader,
}
}
}

/// Extension trait for constructing [`VortexOpenOptions`] from a session.
pub trait OpenOptionsSessionExt:
ArraySessionExt + LayoutSessionExt + RuntimeSessionExt + MemorySessionExt
Expand All @@ -101,7 +87,6 @@ pub trait OpenOptionsSessionExt:
file_size: None,
dtype: None,
footer: None,
initial_read_segments: Default::default(),
metrics_registry: None,
labels: Vec::default(),
cache_layout_reader: false,
Expand Down Expand Up @@ -247,7 +232,7 @@ impl VortexOpenOptions {

let footer = match opts.footer.take() {
Some(footer) => footer,
None => block_on(opts.read_footer(&buffer))?,
None => block_on(opts.read_footer(&buffer))?.footer,
};
footer.validate_file_size(buffer.len() as u64)?;

Expand Down Expand Up @@ -277,18 +262,24 @@ impl VortexOpenOptions {
.clone()
.unwrap_or_else(|| Arc::new(DefaultMetricsRegistry::default()));

let footer = if let Some(footer) = self.footer {
let FooterRead {
footer,
initial_segments,
} = if let Some(footer) = self.footer {
if let Some(file_size) = self.file_size {
footer.validate_file_size(file_size)?;
}
footer
FooterRead {
footer,
initial_segments: HashMap::default(),
}
} else {
self.read_footer(&reader).await?
};

let segment_cache = Arc::new(InstrumentedSegmentCache::new(
InitialReadSegmentCache {
initial: self.initial_read_segments,
initial: initial_segments,
fallback: segment_cache,
},
metrics_registry.as_ref(),
Expand Down Expand Up @@ -319,7 +310,7 @@ impl VortexOpenOptions {
})
}

async fn read_footer(&self, read: &dyn VortexReadAt) -> VortexResult<Footer> {
async fn read_footer(&self, read: &dyn VortexReadAt) -> VortexResult<FooterRead> {
// Fetch the file size and perform the initial read.
let file_size = match self.file_size {
None => read.size().await?,
Expand Down Expand Up @@ -367,23 +358,26 @@ impl VortexOpenOptions {
// If the initial read happened to cover any segments, then we can populate the
// segment cache
let initial_offset = file_size - (deserializer.buffer().len() as u64);
self.populate_initial_segments(initial_offset, deserializer.buffer(), &footer)?;
let initial_segments =
Self::collect_initial_segments(initial_offset, deserializer.buffer(), &footer)?;

Ok(footer)
Ok(FooterRead {
footer,
initial_segments,
})
}

/// Populate segments in the cache that were covered by the initial read.
fn populate_initial_segments(
&self,
/// Collect segments that were covered by the initial read.
fn collect_initial_segments(
initial_offset: u64,
initial_read: &ByteBuffer,
footer: &Footer,
) -> VortexResult<()> {
) -> VortexResult<HashMap<SegmentId, ByteBuffer>> {
let first_idx = footer
.segment_map()
.partition_point(|segment| segment.offset < initial_offset);

let mut initial_read_segments = self.initial_read_segments.write();
let mut initial_read_segments = HashMap::default();

for idx in first_idx..footer.segment_map().len() {
let segment = &footer.segment_map()[idx];
Expand All @@ -410,7 +404,7 @@ impl VortexOpenOptions {
initial_read_segments.insert(segment_id, buffer);
}

Ok(())
Ok(initial_read_segments)
}
}

Expand Down Expand Up @@ -613,14 +607,14 @@ mod tests {
);
}

/// `populate_initial_segments` must bounds-check the segment map against the initial read rather
/// `collect_initial_segments` must bounds-check the segment map against the initial read rather
/// than slicing unchecked, so a segment larger than the read returns an error (see issue #8819).
#[tokio::test]
async fn populate_initial_segments_rejects_out_of_bounds_segment() -> VortexResult<()> {
async fn collect_initial_segments_rejects_out_of_bounds_segment() -> VortexResult<()> {
let session = test_session();

// A valid root layout is obtained by writing and parsing a small file.
// `populate_initial_segments` only consults the segment map, so the layout is irrelevant.
// `collect_initial_segments` only consults the segment map, so the layout is irrelevant.
let mut buf = ByteBufferMut::empty();
let array = Buffer::from((0i32..16).collect::<Vec<i32>>()).into_array();
session
Expand All @@ -630,7 +624,8 @@ mod tests {
let footer = session
.open_options()
.read_footer(&ByteBuffer::from(buf))
.await?;
.await?
.footer;

// Build a footer whose sole segment is far larger than the initial read below.
let bad_segments: Arc<[SegmentSpec]> = Arc::from([SegmentSpec {
Expand All @@ -646,12 +641,9 @@ mod tests {
);

let initial_read = ByteBuffer::zeroed(16);
let Err(err) =
session
.open_options()
.populate_initial_segments(0, &initial_read, &bad_footer)
let Err(err) = VortexOpenOptions::collect_initial_segments(0, &initial_read, &bad_footer)
else {
vortex_bail!("populating an out-of-bounds segment must return an error");
vortex_bail!("collecting an out-of-bounds segment must return an error");
};
assert!(
err.to_string().contains("out of bounds"),
Expand Down
5 changes: 2 additions & 3 deletions vortex-file/src/segments/cache.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
use std::sync::Arc;

use async_trait::async_trait;
use parking_lot::RwLock;
use vortex_buffer::ByteBuffer;
use vortex_error::VortexResult;
use vortex_layout::segments::SegmentCache;
Expand All @@ -14,15 +13,15 @@ use vortex_utils::aliases::hash_map::HashMap;
/// Segment cache containing the initial read segments.
pub struct InitialReadSegmentCache {
/// Segments that were already covered by the footer initial read.
pub initial: RwLock<HashMap<SegmentId, ByteBuffer>>,
pub initial: HashMap<SegmentId, ByteBuffer>,
/// Delegate cache used for all misses and stores.
pub fallback: Arc<dyn SegmentCache>,
}

#[async_trait]
impl SegmentCache for InitialReadSegmentCache {
async fn get(&self, id: SegmentId) -> VortexResult<Option<ByteBuffer>> {
if let Some(buffer) = self.initial.read().get(&id) {
if let Some(buffer) = self.initial.get(&id) {
return Ok(Some(buffer.clone()));
}
self.fallback.get(id).await
Expand Down
Loading