Skip to main content

vortex_file/
open.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use futures::executor::block_on;
7use vortex_array::dtype::DType;
8use vortex_array::memory::MemorySessionExt;
9use vortex_array::session::ArraySessionExt;
10use vortex_buffer::Alignment;
11use vortex_buffer::ByteBuffer;
12use vortex_error::VortexError;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15use vortex_error::vortex_err;
16use vortex_io::VortexReadAt;
17use vortex_io::session::RuntimeSessionExt;
18use vortex_layout::segments::InstrumentedSegmentCache;
19use vortex_layout::segments::NoOpSegmentCache;
20use vortex_layout::segments::SegmentCache;
21use vortex_layout::segments::SegmentCacheSourceAdapter;
22use vortex_layout::segments::SegmentId;
23use vortex_layout::segments::SegmentSource;
24use vortex_layout::segments::SharedSegmentSource;
25use vortex_layout::session::LayoutSessionExt;
26use vortex_metrics::DefaultMetricsRegistry;
27use vortex_metrics::Label;
28use vortex_metrics::MetricsRegistry;
29use vortex_session::VortexSession;
30use vortex_utils::aliases::hash_map::HashMap;
31
32use crate::DeserializeStep;
33use crate::EOF_SIZE;
34use crate::MAX_POSTSCRIPT_SIZE;
35use crate::VortexFile;
36use crate::footer::Footer;
37use crate::segments::BufferSegmentSource;
38use crate::segments::FileSegmentSource;
39use crate::segments::InitialReadSegmentCache;
40use crate::segments::RequestMetrics;
41
42const INITIAL_READ_SIZE: usize = MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE;
43
44struct FooterRead {
45    footer: Footer,
46    initial_segments: HashMap<SegmentId, ByteBuffer>,
47}
48
49/// Open options for a Vortex file reader.
50///
51/// Options are session-bound because opening a file may need array, layout, dtype, runtime, and
52/// memory registries. Construct with [`OpenOptionsSessionExt::open_options`] for the common path.
53///
54/// The opener first resolves a [`Footer`], then creates a segment source for later scans. Known
55/// file metadata can be supplied up front to avoid IO during footer discovery.
56#[derive(Clone)]
57pub struct VortexOpenOptions {
58    /// The session to use for opening the file.
59    session: VortexSession,
60    /// Cache to use for file segments.
61    segment_cache: Option<Arc<dyn SegmentCache>>,
62    /// The number of bytes to read when parsing the footer.
63    initial_read_size: usize,
64    /// An optional, externally provided, file size.
65    file_size: Option<u64>,
66    /// An optional, externally provided, DType.
67    dtype: Option<DType>,
68    /// An optional, externally provided, file layout.
69    footer: Option<Footer>,
70    /// Whether to include user-defined metadata segments when opening the file.
71    include_metadata: bool,
72    /// A metrics registry for the file.
73    metrics_registry: Option<Arc<dyn MetricsRegistry>>,
74    /// Default labels applied to all the file's metrics
75    labels: Vec<Label>,
76    /// Whether to cache file's LayoutReader between scans
77    cache_layout_reader: bool,
78}
79
80/// Extension trait for constructing [`VortexOpenOptions`] from a session.
81pub trait OpenOptionsSessionExt:
82    ArraySessionExt + LayoutSessionExt + RuntimeSessionExt + MemorySessionExt
83{
84    /// Create a new [`VortexOpenOptions`] using the provided session to open a file.
85    fn open_options(&self) -> VortexOpenOptions {
86        VortexOpenOptions {
87            session: self.session(),
88            segment_cache: None,
89            initial_read_size: INITIAL_READ_SIZE,
90            file_size: None,
91            dtype: None,
92            footer: None,
93            include_metadata: false,
94            metrics_registry: None,
95            labels: Vec::default(),
96            cache_layout_reader: false,
97        }
98    }
99}
100impl<S: ArraySessionExt + LayoutSessionExt + RuntimeSessionExt + MemorySessionExt>
101    OpenOptionsSessionExt for S
102{
103}
104
105impl VortexOpenOptions {
106    /// Return the session this opener is bound to.
107    pub fn session(&self) -> &VortexSession {
108        &self.session
109    }
110
111    /// Configure how many bytes to read from the end of the file before parsing the footer.
112    ///
113    /// The actual read is at least large enough to contain the maximum postscript and EOF marker,
114    /// and no larger than the file. Increase this when you expect footer segments to be near the
115    /// end and want to avoid a second footer read.
116    pub fn with_initial_read_size(mut self, initial_read_size: usize) -> Self {
117        self.initial_read_size = initial_read_size;
118        self
119    }
120
121    /// Cache the file's [`LayoutReader`](vortex_layout::LayoutReader) between scans.
122    ///
123    /// This avoids rebuilding the reader tree for repeated scans of the same [`VortexFile`], at the
124    /// cost of keeping reader state alive for the lifetime of the file handle.
125    pub fn with_layout_reader_cache(mut self) -> Self {
126        self.cache_layout_reader = true;
127        self
128    }
129
130    /// Configure a custom [`SegmentCache`].
131    ///
132    /// The cache is checked before the underlying file segment source. Segments covered by the
133    /// initial footer read are also inserted into an internal first-read cache.
134    pub fn with_segment_cache(mut self, segment_cache: Arc<dyn SegmentCache>) -> Self {
135        self.segment_cache = Some(segment_cache);
136        self
137    }
138
139    /// Disable the configured segment cache.
140    ///
141    /// This is useful when deriving an opener for a source whose buffers have different memory
142    /// placement requirements from the configured host cache.
143    pub fn without_segment_cache(mut self) -> Self {
144        self.segment_cache = None;
145        self
146    }
147
148    /// Configure a known file size.
149    ///
150    /// This helps to prevent an I/O request to discover the size of the file.
151    /// Of course, all bets are off if you pass an incorrect value.
152    pub fn with_file_size(mut self, file_size: u64) -> Self {
153        self.file_size = Some(file_size);
154        self
155    }
156
157    /// Configure a known file size.
158    ///
159    /// This helps to prevent an I/O request to discover the size of the file.
160    /// Of course, all bets are off if you pass an incorrect value.
161    pub fn with_some_file_size(mut self, file_size: Option<u64>) -> Self {
162        self.file_size = file_size;
163        self
164    }
165
166    /// Configure a known DType.
167    ///
168    /// If this is provided, then the Vortex file may be opened with fewer I/O requests.
169    ///
170    /// For Vortex files that do not contain a `DType`, this is required.
171    pub fn with_dtype(mut self, dtype: DType) -> Self {
172        self.dtype = Some(dtype);
173        self
174    }
175
176    /// Configure a known file footer.
177    ///
178    /// If this is provided, then the Vortex file can be opened without performing any I/O.
179    /// Once open, the [`Footer`] can be accessed via [`crate::VortexFile::footer`].
180    pub fn with_footer(mut self, footer: Footer) -> Self {
181        self.dtype = Some(footer.layout().dtype().clone());
182        self.footer = Some(footer);
183        self
184    }
185
186    /// Include user-defined metadata segments when opening the file.
187    ///
188    /// By default, opening a file reads only the metadata required to interpret the layout.
189    /// Enabling this option loads all metadata segments named by the postscript, which may require
190    /// additional reads before the footer is returned.
191    pub fn include_metadata(mut self) -> Self {
192        self.include_metadata = true;
193        self
194    }
195
196    /// Configure whether user-defined metadata segments are included when opening the file.
197    ///
198    /// Enabling this option loads all metadata segments named by the postscript, which may require
199    /// additional reads before the footer is returned.
200    pub fn with_include_metadata(mut self, include_metadata: bool) -> Self {
201        self.include_metadata = include_metadata;
202        self
203    }
204
205    /// Configure a custom [`MetricsRegistry`] implementation.
206    pub fn with_metrics_registry(mut self, metrics: Arc<dyn MetricsRegistry>) -> Self {
207        self.metrics_registry = Some(metrics);
208        self
209    }
210
211    /// Adds labels to all the file's metrics.
212    pub fn with_labels(mut self, labels: Vec<Label>) -> Self {
213        self.labels.extend(labels);
214        self
215    }
216
217    /// Open a Vortex file using the provided I/O source.
218    ///
219    /// This is the most common way to open a [`VortexFile`] and tends to provide the best
220    /// out-of-the-box performance. The underlying I/O system will continue to be optimised for
221    /// different file systems and object stores so we encourage users to use this method
222    /// whenever possible and file issues if they encounter problems.
223    pub async fn open(self, source: Arc<dyn VortexReadAt>) -> VortexResult<VortexFile> {
224        self.open_read(source).await
225    }
226
227    /// Open a Vortex file from a filesystem path.
228    #[cfg(not(target_arch = "wasm32"))]
229    pub async fn open_path(self, path: impl AsRef<std::path::Path>) -> VortexResult<VortexFile> {
230        use vortex_io::std_file::FileReadAt;
231        let handle = self.session.handle();
232        let allocator = self.session.allocator();
233        let source = Arc::new(FileReadAt::open_with_allocator(path, handle, allocator)?);
234        self.open(source).await
235    }
236
237    /// Open a Vortex file from an in-memory buffer.
238    ///
239    /// This uses a `BufferSegmentSource` that resolves segments synchronously
240    /// by slicing the buffer directly, bypassing the async I/O pipeline.
241    ///
242    /// Segment cache and metrics registry settings are ignored for this path.
243    pub fn open_buffer<B: Into<ByteBuffer>>(self, buffer: B) -> VortexResult<VortexFile> {
244        let buffer: ByteBuffer = buffer.into();
245
246        if self.segment_cache.is_some() {
247            tracing::warn!("segment cache is ignored for in-memory `open_buffer`");
248        }
249        if self.metrics_registry.is_some() {
250            tracing::warn!("metrics registry is ignored for in-memory `open_buffer`");
251        }
252
253        let cache_layout_reader = self.cache_layout_reader;
254        let include_metadata = self.include_metadata;
255        let mut opts = self.with_initial_read_size(0);
256
257        let footer = match opts.footer.take() {
258            Some(footer) => footer,
259            None => block_on(opts.read_footer(&buffer))?.footer,
260        };
261        footer.validate_file_size(buffer.len() as u64)?;
262
263        let segment_source: Arc<dyn SegmentSource> = Arc::new(BufferSegmentSource::new(
264            buffer,
265            footer.segment_specs_with_metadata(),
266        ));
267        let metadata = if include_metadata {
268            block_on(resolve_metadata(&footer, Arc::clone(&segment_source)))?
269        } else {
270            Arc::new(HashMap::new())
271        };
272        let file = VortexFile::new(footer, segment_source, opts.session).with_metadata(metadata);
273        Ok(if cache_layout_reader {
274            file.with_caching()
275        } else {
276            file
277        })
278    }
279
280    /// Open a [`VortexFile`] using any [`VortexReadAt`] implementation.
281    ///
282    /// This is the common path for files, object stores, and custom random-access sources.
283    pub async fn open_read<R: VortexReadAt + Clone>(self, reader: R) -> VortexResult<VortexFile> {
284        let segment_cache = self
285            .segment_cache
286            .clone()
287            .unwrap_or_else(|| Arc::new(NoOpSegmentCache));
288
289        let metrics_registry = self
290            .metrics_registry
291            .clone()
292            .unwrap_or_else(|| Arc::new(DefaultMetricsRegistry::default()));
293
294        let FooterRead {
295            footer,
296            initial_segments,
297        } = if let Some(footer) = self.footer {
298            if let Some(file_size) = self.file_size {
299                footer.validate_file_size(file_size)?;
300            }
301            FooterRead {
302                footer,
303                initial_segments: HashMap::default(),
304            }
305        } else {
306            self.read_footer(&reader).await?
307        };
308
309        // With a segment cache, also store the segments that the footer read covered, so that a
310        // later open given this footer finds them there, instead of reading them again. Keep them
311        // in the initial map too: the cache can evict them, or not admit them at all, and the
312        // buffers are shared, not copied.
313        if self.segment_cache.is_some() {
314            for (id, buffer) in &initial_segments {
315                segment_cache.put(*id, buffer.clone()).await?;
316            }
317        }
318
319        let segment_cache = Arc::new(InstrumentedSegmentCache::new(
320            InitialReadSegmentCache {
321                initial: initial_segments,
322                fallback: segment_cache,
323            },
324            metrics_registry.as_ref(),
325            self.labels.clone(),
326        ));
327
328        let metrics = RequestMetrics::new(metrics_registry.as_ref(), self.labels);
329
330        // Create a segment source backed by the VortexRead implementation.
331        let segment_source = Arc::new(SharedSegmentSource::new(FileSegmentSource::open(
332            footer.segment_specs_with_metadata(),
333            reader,
334            self.session.handle(),
335            metrics,
336        )));
337
338        // Wrap up the segment source to first resolve segments from the initial read cache.
339        let segment_source: Arc<dyn SegmentSource> = Arc::new(SegmentCacheSourceAdapter::new(
340            segment_cache,
341            segment_source,
342        ));
343
344        let metadata = if self.include_metadata {
345            resolve_metadata(&footer, Arc::clone(&segment_source)).await?
346        } else {
347            Arc::new(HashMap::new())
348        };
349        let file =
350            VortexFile::new(footer, segment_source, self.session.clone()).with_metadata(metadata);
351        Ok(if self.cache_layout_reader {
352            file.with_caching()
353        } else {
354            file
355        })
356    }
357
358    async fn read_footer(&self, read: &dyn VortexReadAt) -> VortexResult<FooterRead> {
359        // Fetch the file size and perform the initial read.
360        let file_size = match self.file_size {
361            None => read.size().await?,
362            Some(file_size) => file_size,
363        };
364        let mut initial_read_size = self
365            .initial_read_size
366            // Make sure we read enough to cover the postscript
367            .max(MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE);
368        if let Ok(file_size) = usize::try_from(file_size) {
369            initial_read_size = initial_read_size.min(file_size);
370        }
371
372        let initial_offset = file_size - initial_read_size as u64;
373        let initial_read: ByteBuffer = read
374            .read_at(initial_offset, initial_read_size, Alignment::none())
375            .await?
376            .try_into_host()?
377            .await?;
378
379        let mut deserializer = Footer::deserializer(initial_read, self.session.clone())
380            .with_size(file_size)
381            .with_some_dtype(self.dtype.clone());
382
383        let footer = loop {
384            match deserializer.deserialize()? {
385                DeserializeStep::NeedMoreData { offset, len } => {
386                    let more_data = read
387                        .read_at(offset, len, Alignment::none())
388                        .await?
389                        .try_into_host()?
390                        .await?;
391                    deserializer.prefix_data(more_data);
392                }
393                DeserializeStep::NeedFileSize => unreachable!("We passed file_size above"),
394                DeserializeStep::Done(footer) => break Ok::<_, VortexError>(footer),
395            }
396        }?;
397
398        // Segment specs describe the data file, not necessarily the byte stream used to
399        // deserialize a standalone cached footer. Validate them here, where we know this is the
400        // size of the actual data source (see issue #8819).
401        footer.validate_file_size(file_size)?;
402
403        // If the initial read happened to cover any segments, then we can populate the
404        // segment cache
405        let initial_offset = file_size - (deserializer.buffer().len() as u64);
406        let initial_segments =
407            Self::collect_initial_segments(initial_offset, deserializer.buffer(), &footer)?;
408
409        Ok(FooterRead {
410            footer,
411            initial_segments,
412        })
413    }
414
415    /// Collect segments that were covered by the initial read.
416    fn collect_initial_segments(
417        initial_offset: u64,
418        initial_read: &ByteBuffer,
419        footer: &Footer,
420    ) -> VortexResult<HashMap<SegmentId, ByteBuffer>> {
421        let mut initial_read_segments = HashMap::default();
422
423        // Iterate `segment_specs_with_metadata` (not just the segment map) so metadata segments
424        // covered by the initial read are cached too. Metadata segments are appended and not
425        // offset-sorted, so we skip per-segment rather than partition on offset.
426        for (idx, segment) in footer.segment_specs_with_metadata().iter().enumerate() {
427            if segment.offset < initial_offset {
428                continue;
429            }
430            let segment_id =
431                SegmentId::from(u32::try_from(idx).vortex_expect("Invalid segment ID"));
432            let offset =
433                usize::try_from(segment.offset - initial_offset).vortex_expect("Invalid offset");
434            // The segment map is validated against the file size before this method is called, but
435            // still bounds-check here so slicing never depends on that distant validation for panic
436            // safety (see issue #8819).
437            let end = offset
438                .checked_add(segment.length as usize)
439                .filter(|end| *end <= initial_read.len())
440                .ok_or_else(|| {
441                    vortex_err!(
442                        "Segment at offset {} with length {} is out of bounds of the \
443                         {}-byte initial read",
444                        segment.offset,
445                        segment.length,
446                        initial_read.len(),
447                    )
448                })?;
449            let buffer = initial_read.slice(offset..end).aligned(segment.alignment);
450            initial_read_segments.insert(segment_id, buffer);
451        }
452
453        Ok(initial_read_segments)
454    }
455}
456
457async fn resolve_metadata(
458    footer: &Footer,
459    segment_source: Arc<dyn SegmentSource>,
460) -> VortexResult<Arc<HashMap<String, ByteBuffer>>> {
461    let first_metadata_id = footer.segment_map().len();
462    let requests = footer
463        .metadata_segments()
464        .enumerate()
465        .map(|(index, (key, locator))| {
466            let id = u32::try_from(first_metadata_id + index).map(SegmentId::from);
467            let key = key.to_string();
468            let alignment = locator.alignment;
469            let segment_source = Arc::clone(&segment_source);
470            async move {
471                let handle = segment_source.request(id?).await?;
472                let buffer = handle.try_into_host()?.await?;
473                Ok::<_, VortexError>((
474                    key,
475                    ByteBuffer::copy_from_aligned(buffer.as_slice(), alignment),
476                ))
477            }
478        });
479    let metadata = futures::future::try_join_all(requests)
480        .await?
481        .into_iter()
482        .collect::<HashMap<_, _>>();
483
484    Ok(Arc::new(metadata))
485}
486
487#[cfg(feature = "object_store")]
488impl VortexOpenOptions {
489    /// Open a Vortex file from an `object_store` backend and path.
490    ///
491    /// `path` is the object's *literal* key, exactly as the store holds it.
492    pub async fn open_object_store(
493        self,
494        object_store: &Arc<dyn object_store::ObjectStore>,
495        path: object_store::path::Path,
496    ) -> VortexResult<VortexFile> {
497        use vortex_io::object_store::ObjectStoreReadAt;
498
499        let handle = self.session.handle();
500        let allocator = self.session.allocator();
501        let source = Arc::new(ObjectStoreReadAt::new_with_allocator(
502            Arc::clone(object_store),
503            path,
504            handle,
505            allocator,
506        ));
507        self.open(source).await
508    }
509}
510
511#[cfg(test)]
512mod tests {
513    use std::alloc::Layout;
514    use std::ptr::NonNull;
515    use std::sync::atomic::AtomicUsize;
516    use std::sync::atomic::Ordering;
517
518    use allocator_api2::alloc::AllocError;
519    use allocator_api2::alloc::Allocator;
520    use allocator_api2::alloc::Global;
521    use futures::future::BoxFuture;
522    use parking_lot::Mutex;
523    use vortex_array::IntoArray;
524    use vortex_array::buffer::BufferHandle;
525    use vortex_array::memory::BufferAllocatorRef;
526    use vortex_array::memory::MemorySessionExt;
527    use vortex_buffer::Alignment;
528    use vortex_buffer::Buffer;
529    use vortex_buffer::ByteBuffer;
530    use vortex_buffer::ByteBufferMut;
531    use vortex_error::vortex_bail;
532    use vortex_io::session::RuntimeSession;
533    use vortex_layout::session::LayoutSession;
534    use vortex_session::registry::Id;
535    use vortex_session::registry::ReadContext;
536
537    use super::*;
538    use crate::WriteOptionsSessionExt;
539    use crate::footer::SegmentSpec;
540
541    fn test_session() -> VortexSession {
542        let session = vortex_array::array_session()
543            .with::<LayoutSession>()
544            .with::<RuntimeSession>();
545        crate::register_default_encodings(&session);
546        crate::enable_all_registered_array_encodings(&session);
547        session
548    }
549
550    #[derive(Clone)]
551    // Define CountingRead struct
552    struct CountingRead<R> {
553        inner: R,
554        total_read: Arc<AtomicUsize>,
555        first_read_len: Arc<AtomicUsize>,
556        reads: Arc<Mutex<Vec<(u64, usize)>>>,
557    }
558
559    impl<R: VortexReadAt + Clone> VortexReadAt for CountingRead<R> {
560        fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
561            self.inner.size()
562        }
563
564        fn read_at(
565            &self,
566            offset: u64,
567            length: usize,
568            alignment: Alignment,
569        ) -> BoxFuture<'static, VortexResult<BufferHandle>> {
570            self.total_read.fetch_add(length, Ordering::Relaxed);
571            self.reads.lock().push((offset, length));
572            let _ = self.first_read_len.compare_exchange(
573                0,
574                length,
575                Ordering::Relaxed,
576                Ordering::Relaxed,
577            );
578            self.inner.read_at(offset, length, alignment)
579        }
580
581        fn concurrency(&self) -> usize {
582            self.inner.concurrency()
583        }
584    }
585
586    #[derive(Debug)]
587    struct CountingAllocator {
588        allocations: Arc<AtomicUsize>,
589    }
590
591    // SAFETY: this forwards memory operations to Global and only counts allocations.
592    unsafe impl Allocator for CountingAllocator {
593        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
594            self.allocations.fetch_add(1, Ordering::Relaxed);
595            Global.allocate(layout)
596        }
597
598        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
599            // SAFETY: ptr and layout came from Global.
600            unsafe { Global.deallocate(ptr, layout) }
601        }
602    }
603
604    #[tokio::test]
605    async fn test_initial_read_size() {
606        let session = vortex_array::array_session()
607            .with::<LayoutSession>()
608            .with::<RuntimeSession>();
609
610        crate::register_default_encodings(&session);
611        crate::enable_all_registered_array_encodings(&session);
612
613        // Create a large file (> 1MB)
614        let mut buf = ByteBufferMut::empty();
615
616        // 1.5M integers -> ~6MB. We use high-entropy (pseudo-random) values so the data does not
617        // compress well under any encoding (Sequence, RunEnd, Delta, ...), keeping the written
618        // file comfortably above 1MB.
619        let mut state = 0x9E37_79B9u32;
620        let array = Buffer::from(
621            (0i32..1_500_000)
622                .map(|_| {
623                    state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
624                    state as i32
625                })
626                .collect::<Vec<i32>>(),
627        )
628        .into_array();
629
630        session
631            .write_options()
632            .write(&mut buf, array.to_array_stream())
633            .await
634            .unwrap();
635
636        let buffer = ByteBuffer::from(buf);
637        assert!(
638            buffer.len() > 1024 * 1024,
639            "Buffer length is only {} bytes",
640            buffer.len()
641        );
642
643        let total_read = Arc::new(AtomicUsize::new(0));
644        let first_read_len = Arc::new(AtomicUsize::new(0));
645        let reader = CountingRead {
646            inner: buffer,
647            total_read: Arc::clone(&total_read),
648            first_read_len: Arc::clone(&first_read_len),
649            reads: Arc::new(Mutex::new(Vec::new())),
650        };
651
652        // Open the file
653        let _file = session.open_options().open_read(reader).await.unwrap();
654
655        // Assert that we read approximately the postscript size, not 1MB
656        let first = first_read_len.load(Ordering::Relaxed);
657        assert_eq!(
658            first,
659            MAX_POSTSCRIPT_SIZE as usize + EOF_SIZE,
660            "Read exactly the postscript size"
661        );
662        let read = total_read.load(Ordering::Relaxed);
663        assert!(read < 1024 * 1024, "Read {} bytes, expected < 1MB", read);
664    }
665
666    #[tokio::test]
667    async fn test_metadata_outside_initial_read_uses_targeted_read() -> VortexResult<()> {
668        let session = vortex_array::array_session()
669            .with::<LayoutSession>()
670            .with::<RuntimeSession>();
671        crate::register_default_encodings(&session);
672        crate::enable_all_registered_array_encodings(&session);
673
674        let metadata = ByteBuffer::copy_from(vec![0x5a; INITIAL_READ_SIZE * 2]);
675        let mut output = ByteBufferMut::empty();
676        let summary = session
677            .write_options()
678            .with_metadata_segment("outside", metadata.clone())
679            .write(
680                &mut output,
681                Buffer::from(vec![1u32]).into_array().to_array_stream(),
682            )
683            .await?;
684        let locator = *summary
685            .footer()
686            .metadata_segment("outside")
687            .vortex_expect("metadata locator");
688        let bytes = ByteBuffer::from(output);
689        assert!(locator.offset < bytes.len() as u64 - INITIAL_READ_SIZE as u64);
690
691        let default_total = Arc::new(AtomicUsize::new(0));
692        let default_reads = Arc::new(Mutex::new(Vec::new()));
693        let default_reader = CountingRead {
694            inner: bytes.clone(),
695            total_read: Arc::clone(&default_total),
696            first_read_len: Arc::new(AtomicUsize::new(0)),
697            reads: Arc::clone(&default_reads),
698        };
699        let default_file = session.open_options().open_read(default_reader).await?;
700        assert!(default_file.metadata_segment("outside").is_none());
701        assert!(
702            !default_reads
703                .lock()
704                .contains(&(locator.offset, locator.length as usize))
705        );
706        // A default open must not amplify into the metadata: total bytes read stay
707        // below the metadata segment's size (itself 2x the initial read).
708        assert!(
709            default_total.load(Ordering::Relaxed) < metadata.len(),
710            "default open read {} bytes; metadata segment is {}",
711            default_total.load(Ordering::Relaxed),
712            metadata.len()
713        );
714
715        let metadata_reads = Arc::new(Mutex::new(Vec::new()));
716        let metadata_reader = CountingRead {
717            inner: bytes,
718            total_read: Arc::new(AtomicUsize::new(0)),
719            first_read_len: Arc::new(AtomicUsize::new(0)),
720            reads: Arc::clone(&metadata_reads),
721        };
722        let file = session
723            .open_options()
724            .include_metadata()
725            .open_read(metadata_reader)
726            .await?;
727        assert_eq!(
728            file.metadata_segment("outside").map(ByteBuffer::as_slice),
729            Some(metadata.as_slice())
730        );
731        assert!(
732            metadata_reads
733                .lock()
734                .contains(&(locator.offset, locator.length as usize))
735        );
736
737        Ok(())
738    }
739
740    #[tokio::test]
741    async fn test_with_footer_include_metadata_open_buffer_resolves() -> VortexResult<()> {
742        let session = vortex_array::array_session()
743            .with::<LayoutSession>()
744            .with::<RuntimeSession>();
745        crate::register_default_encodings(&session);
746        crate::enable_all_registered_array_encodings(&session);
747
748        let value = ByteBuffer::copy_from(b"supplied-footer metadata");
749        let mut output = ByteBufferMut::empty();
750        let summary = session
751            .write_options()
752            .with_metadata_segment("key", value.clone())
753            .write(
754                &mut output,
755                Buffer::from(vec![1u32]).into_array().to_array_stream(),
756            )
757            .await?;
758        let footer = summary.footer().clone();
759        let bytes = ByteBuffer::from(output);
760
761        let file = session
762            .open_options()
763            .with_footer(footer.clone())
764            .include_metadata()
765            .open_buffer(bytes.clone())?;
766        assert_eq!(
767            file.metadata_segment("key").map(ByteBuffer::as_slice),
768            Some(value.as_slice())
769        );
770
771        let default = session
772            .open_options()
773            .with_footer(footer)
774            .open_buffer(bytes)?;
775        assert!(default.metadata_segment("key").is_none());
776
777        Ok(())
778    }
779
780    #[cfg(not(target_arch = "wasm32"))]
781    #[tokio::test]
782    async fn test_open_path_uses_memory_session_allocator() {
783        let session = vortex_array::array_session()
784            .with::<LayoutSession>()
785            .with::<RuntimeSession>();
786
787        crate::register_default_encodings(&session);
788        crate::enable_all_registered_array_encodings(&session);
789
790        let mut buf = ByteBufferMut::empty();
791        let array = Buffer::from((0i32..16_384).collect::<Vec<i32>>()).into_array();
792        session
793            .write_options()
794            .write(&mut buf, array.to_array_stream())
795            .await
796            .unwrap();
797
798        let file_path = std::env::temp_dir().join(format!(
799            "vortex-open-memory-session-{}.vx",
800            std::process::id()
801        ));
802        std::fs::write(&file_path, ByteBuffer::from(buf).as_slice()).unwrap();
803
804        let allocations = Arc::new(AtomicUsize::new(0));
805        let session = session.with_allocator(BufferAllocatorRef::new(CountingAllocator {
806            allocations: Arc::clone(&allocations),
807        }));
808
809        let _file = session.open_options().open_path(&file_path).await.unwrap();
810        std::fs::remove_file(&file_path).unwrap();
811
812        assert!(
813            allocations.load(Ordering::Relaxed) > 0,
814            "expected at least one host allocation from MemorySession"
815        );
816    }
817
818    /// `collect_initial_segments` must bounds-check the segment map against the initial read rather
819    /// than slicing unchecked, so a segment larger than the read returns an error (see issue #8819).
820    #[tokio::test]
821    async fn collect_initial_segments_rejects_out_of_bounds_segment() -> VortexResult<()> {
822        let session = test_session();
823
824        // A valid root layout is obtained by writing and parsing a small file.
825        // `collect_initial_segments` only consults the segment map, so the layout is irrelevant.
826        let mut buf = ByteBufferMut::empty();
827        let array = Buffer::from((0i32..16).collect::<Vec<i32>>()).into_array();
828        session
829            .write_options()
830            .write(&mut buf, array.to_array_stream())
831            .await?;
832        let footer = session
833            .open_options()
834            .read_footer(&ByteBuffer::from(buf))
835            .await?
836            .footer;
837
838        // Build a footer whose sole segment is far larger than the initial read below.
839        let bad_segments: Arc<[SegmentSpec]> = Arc::from([SegmentSpec {
840            offset: 0,
841            length: 1024,
842            alignment: Alignment::none(),
843        }]);
844        let bad_footer = Footer::new(
845            Arc::clone(footer.layout()),
846            bad_segments,
847            None,
848            ReadContext::new(Vec::<Id>::new()),
849        );
850
851        let initial_read = ByteBuffer::zeroed(16);
852        let Err(err) = VortexOpenOptions::collect_initial_segments(0, &initial_read, &bad_footer)
853        else {
854            vortex_bail!("collecting an out-of-bounds segment must return an error");
855        };
856        assert!(
857            err.to_string().contains("out of bounds"),
858            "unexpected error: {err}"
859        );
860
861        Ok(())
862    }
863
864    #[tokio::test]
865    async fn standalone_footer_round_trip() -> VortexResult<()> {
866        let session = test_session();
867
868        let mut file_bytes = ByteBufferMut::empty();
869        let values = (0i32..65_536)
870            .map(|value| value.wrapping_mul(1_664_525).wrapping_add(1_013_904_223))
871            .collect::<Vec<_>>();
872        let array = Buffer::from(values).into_array();
873        session
874            .write_options()
875            .write(&mut file_bytes, array.to_array_stream())
876            .await?;
877        let file_bytes = ByteBuffer::from(file_bytes);
878
879        let footer = session
880            .open_options()
881            .open_buffer(file_bytes.clone())?
882            .footer()
883            .clone();
884        let last_segment_end = footer
885            .segment_map()
886            .iter()
887            .map(|segment| segment.offset + u64::from(segment.length))
888            .max()
889            .unwrap_or_default();
890        let serialized_footer = footer.into_serializer().serialize()?;
891        let serialized_footer_size = serialized_footer.iter().map(ByteBuffer::len).sum();
892        assert!(last_segment_end > serialized_footer_size as u64);
893        let mut footer_bytes = ByteBufferMut::with_capacity(serialized_footer_size);
894        for buffer in serialized_footer {
895            footer_bytes.extend_from_slice(&buffer);
896        }
897
898        let mut deserializer = Footer::deserializer(footer_bytes.freeze(), session.clone())
899            .with_size(serialized_footer_size as u64);
900        let DeserializeStep::Done(cached_footer) = deserializer.deserialize()? else {
901            vortex_bail!("standalone footer bytes must be sufficient for deserialization");
902        };
903
904        session
905            .open_options()
906            .with_footer(cached_footer)
907            .open_buffer(file_bytes)?;
908
909        Ok(())
910    }
911}