vortex_file/multi/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Builder for constructing a [`MultiLayoutDataSource`] from multiple Vortex files.
5
6mod session;
7
8use std::sync::Arc;
9
10use async_trait::async_trait;
11use futures::TryStreamExt;
12use session::MultiFileSessionExt;
13use tracing::debug;
14use vortex_error::VortexResult;
15use vortex_error::vortex_bail;
16use vortex_io::filesystem::FileListing;
17use vortex_io::filesystem::FileSystemRef;
18use vortex_layout::LayoutReaderRef;
19use vortex_layout::scan::multi::LayoutReaderFactory;
20use vortex_layout::scan::multi::MultiLayoutDataSource;
21use vortex_scan::DataSource;
22use vortex_session::VortexSession;
23
24use crate::OpenOptionsSessionExt;
25use crate::VortexOpenOptions;
26use crate::v2::FileStatsLayoutReader;
27
28/// A builder that discovers multiple Vortex files from glob patterns and constructs a
29/// [`MultiLayoutDataSource`] to scan them as a single data source.
30///
31/// The primary interface is [`Self::with_glob`], which accepts a glob pattern and an optional
32/// filesystem. For non-local filesystems (S3, GCS, etc.), callers must provide a [`FileSystemRef`].
33/// For local files, pass `None` and a local filesystem will be created automatically.
34///
35/// # Examples
36///
37/// ```ignore
38/// // Local files — filesystem is auto-created:
39/// let ds = MultiFileDataSource::new(session)
40///     .with_glob("/data/warehouse/*.vortex", None)
41///     .build()
42///     .await?;
43///
44/// // S3 — caller provides the filesystem:
45/// let ds = MultiFileDataSource::new(session)
46///     .with_glob("prefix/*.vortex", Some(s3_fs))
47///     .build()
48///     .await?;
49///
50/// // Mixed filesystems — multiple globs with different filesystems:
51/// let ds = MultiFileDataSource::new(session)
52///     .with_glob("bucket-a/*.vortex", Some(s3_fs.clone()))
53///     .with_glob("bucket-b/*.vortex", Some(s3_fs))
54///     .with_glob("gcs-bucket/*.vortex", Some(gcs_fs))
55///     .build()
56///     .await?;
57/// ```
58pub struct MultiFileDataSource {
59    session: VortexSession,
60    /// List of (glob, optional filesystem) pairs to resolve.
61    /// When the filesystem is None, a local filesystem will be created in build().
62    glob_sources: Vec<(String, Option<FileSystemRef>)>,
63    open_options_fn: Arc<dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync>,
64}
65
66impl MultiFileDataSource {
67    /// Create a new [`MultiFileDataSource`] builder.
68    pub fn new(session: VortexSession) -> Self {
69        Self {
70            session,
71            glob_sources: Vec::new(),
72            open_options_fn: Arc::new(|opts| opts),
73        }
74    }
75
76    /// Add a path glob for file discovery.
77    ///
78    /// The glob path should be relative to the filesystem's base URL. Pass `None` for the
79    /// filesystem to use the local filesystem (auto-created in [`Self::build`]).
80    pub fn with_glob(mut self, glob: impl Into<String>, fs: Option<FileSystemRef>) -> Self {
81        let glob_str = glob.into().trim_start_matches('/').to_string();
82        self.glob_sources.push((glob_str, fs));
83        self
84    }
85
86    /// Customize [`VortexOpenOptions`] applied to each file.
87    ///
88    /// Use this to configure segment caches, metrics registries, or other per-file options.
89    pub fn with_open_options(
90        mut self,
91        f: impl Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync + 'static,
92    ) -> Self {
93        self.open_options_fn = Arc::new(f);
94        self
95    }
96
97    /// Build the [`DataSource`].
98    ///
99    /// Discovers files via glob, opens the first file eagerly to determine the schema,
100    /// and creates lazy factories for the remaining files.
101    pub async fn build(self) -> VortexResult<impl DataSource> {
102        if self.glob_sources.is_empty() {
103            vortex_bail!("MultiFileDataSource requires at least one glob pattern");
104        }
105
106        // Create local filesystem lazily if needed (only if any glob lacks a filesystem).
107        let local_fs: Option<FileSystemRef> = self
108            .glob_sources
109            .iter()
110            .any(|(_, fs)| fs.is_none())
111            .then(|| create_local_filesystem(&self.session))
112            .transpose()?;
113
114        // Collect files from all glob sources.
115        let mut all_files: Vec<(FileListing, FileSystemRef)> = Vec::new();
116        for (glob, maybe_fs) in &self.glob_sources {
117            // Use the provided filesystem, or fall back to the local filesystem.
118            // We know local_fs is Some when maybe_fs is None (by construction above).
119            let fs = maybe_fs
120                .as_ref()
121                .or(local_fs.as_ref())
122                .map(Arc::clone)
123                .unwrap_or_else(|| {
124                    unreachable!("local_fs is set when any glob lacks a filesystem")
125                });
126            let files: Vec<FileListing> = fs.glob(glob)?.try_collect().await?;
127            for file in files {
128                all_files.push((file, Arc::clone(&fs)));
129            }
130        }
131
132        if all_files.is_empty() {
133            let globs: Vec<_> = self.glob_sources.iter().map(|(g, _)| g.as_str()).collect();
134            vortex_bail!("No files matched the glob pattern(s): {:?}", globs);
135        }
136
137        let file_count = all_files.len();
138        let globs: Vec<_> = self.glob_sources.iter().map(|(g, _)| g.as_str()).collect();
139        debug!(file_count, glob = ?globs, "discovered files");
140
141        // Open first file eagerly for dtype.
142        let (first_file_listing, first_fs) = &all_files[0];
143        let first_file = open_file(
144            first_fs,
145            first_file_listing,
146            &self.session,
147            self.open_options_fn.as_ref(),
148        )
149        .await?;
150        let first_reader = layout_reader_with_stats(&first_file)?;
151
152        let factories: Vec<Arc<dyn LayoutReaderFactory>> = all_files[1..]
153            .iter()
154            .map(|(file, fs)| {
155                Arc::new(VortexFileReaderFactory {
156                    fs: Arc::clone(fs),
157                    file: file.clone(),
158                    session: self.session.clone(),
159                    open_options_fn: Arc::clone(&self.open_options_fn),
160                }) as Arc<dyn LayoutReaderFactory>
161            })
162            .collect();
163
164        let inner = MultiLayoutDataSource::new_with_first(first_reader, factories, &self.session);
165
166        debug!(file_count, dtype = %inner.dtype(), "built MultiFileDataSource");
167
168        Ok(inner)
169    }
170}
171
172/// Creates a local filesystem backed by `object_store::local::LocalFileSystem`.
173// TODO(ngates): create a native file system without an object_store dependency.
174//  Turns out it's not a trivial change because we have always used object_store with its own
175//  coalescing and concurrency configs, so we need to re-tune for local disk.
176#[cfg(feature = "object_store")]
177fn create_local_filesystem(session: &VortexSession) -> VortexResult<FileSystemRef> {
178    use vortex_io::object_store::ObjectStoreFileSystem;
179    use vortex_io::session::RuntimeSessionExt;
180
181    let store = Arc::new(object_store::local::LocalFileSystem::default());
182    let fs: FileSystemRef = Arc::new(ObjectStoreFileSystem::new(store, session.handle()));
183    Ok(fs)
184}
185
186#[cfg(not(feature = "object_store"))]
187fn create_local_filesystem(_session: &VortexSession) -> VortexResult<FileSystemRef> {
188    vortex_bail!(
189        "The 'object_store' feature is required for automatic local filesystem creation. \
190             Either enable the feature or provide a filesystem via .with_filesystem()."
191    );
192}
193
194/// Open a single Vortex file, checking the session's footer cache.
195async fn open_file(
196    fs: &FileSystemRef,
197    file: &FileListing,
198    session: &VortexSession,
199    open_options_fn: &(dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync),
200) -> VortexResult<crate::VortexFile> {
201    debug!(path = %file.path, "opening vortex file");
202
203    // Open the reader first so we can use its URI as the cache key.
204    // The URI includes the full path (with any filesystem prefix), making it unique
205    // even when different PrefixFileSystem instances strip paths to the same relative name.
206    let source = fs.open_read(&file.path).await?;
207    let cache_key = source
208        .uri()
209        .map(|u| u.to_string())
210        .unwrap_or_else(|| file.path.clone());
211
212    // Build open options. The DashMap Ref from multi_file() must not live across an await,
213    // so we scope the cache lookup in a block.
214    let options = {
215        let mut options = open_options_fn(session.open_options());
216        if let Some(size) = file.size {
217            options = options.with_file_size(size);
218        }
219        if let Some(footer) = session.multi_file().get_footer(&cache_key) {
220            options = options.with_footer(footer);
221        }
222        options
223    };
224
225    let vortex_file = options.open(source).await?;
226
227    // Store footer in cache (scoped to avoid holding the Ref across subsequent code).
228    session
229        .multi_file()
230        .put_footer(&cache_key, vortex_file.footer().clone());
231    Ok(vortex_file)
232}
233
234/// Creates a layout reader from a VortexFile, wrapping with `FileStatsLayoutReader` when
235/// file-level statistics are available.
236fn layout_reader_with_stats(file: &crate::VortexFile) -> VortexResult<LayoutReaderRef> {
237    let mut reader = file.layout_reader()?;
238    if let Some(stats) = file.file_stats().cloned() {
239        reader = Arc::new(FileStatsLayoutReader::new(
240            reader,
241            stats,
242            file.session.clone(),
243        ));
244    }
245    Ok(reader)
246}
247
248/// A [`LayoutReaderFactory`] that lazily opens a single Vortex file and returns its layout reader.
249struct VortexFileReaderFactory {
250    fs: FileSystemRef,
251    file: FileListing,
252    session: VortexSession,
253    open_options_fn: Arc<dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync>,
254}
255
256#[async_trait]
257impl LayoutReaderFactory for VortexFileReaderFactory {
258    async fn open(&self) -> VortexResult<Option<LayoutReaderRef>> {
259        let file = open_file(
260            &self.fs,
261            &self.file,
262            &self.session,
263            self.open_options_fn.as_ref(),
264        )
265        .await?;
266        Ok(Some(layout_reader_with_stats(&file)?))
267    }
268}
vortex_file/multi/mod.rs

vortex_file/multi/
mod.rs
