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    ///
81    /// Relative paths are resolved against the process working directory.
82    pub fn with_glob(mut self, glob: impl Into<String>, fs: Option<FileSystemRef>) -> Self {
83        let glob = glob.into();
84        let glob = if fs.is_none() && std::path::Path::new(&glob).is_relative() {
85            std::env::current_dir()
86                .map(|cwd| cwd.join(&glob).to_string_lossy().into_owned())
87                .unwrap_or(glob)
88                .trim_start_matches('/')
89                .to_string()
90        } else {
91            glob.trim_start_matches('/').to_string()
92        };
93        self.glob_sources.push((glob, fs));
94        self
95    }
96
97    /// Customize [`VortexOpenOptions`] applied to each file.
98    ///
99    /// Use this to configure segment caches, metrics registries, or other per-file options.
100    pub fn with_open_options(
101        mut self,
102        f: impl Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync + 'static,
103    ) -> Self {
104        self.open_options_fn = Arc::new(f);
105        self
106    }
107
108    /// Build the [`DataSource`].
109    ///
110    /// Discovers files via glob, opens the first file eagerly to determine the schema,
111    /// and creates lazy factories for the remaining files.
112    pub async fn build(self) -> VortexResult<MultiLayoutDataSource> {
113        if self.glob_sources.is_empty() {
114            vortex_bail!("MultiFileDataSource requires at least one glob pattern");
115        }
116
117        // Create local filesystem lazily if needed (only if any glob lacks a filesystem).
118        let local_fs: Option<FileSystemRef> = self
119            .glob_sources
120            .iter()
121            .any(|(_, fs)| fs.is_none())
122            .then(|| create_local_filesystem(&self.session))
123            .transpose()?;
124
125        // Collect files from all glob sources.
126        let mut all_files: Vec<(FileListing, FileSystemRef)> = Vec::new();
127        for (glob, maybe_fs) in &self.glob_sources {
128            // Use the provided filesystem, or fall back to the local filesystem.
129            // We know local_fs is Some when maybe_fs is None (by construction above).
130            let fs = maybe_fs
131                .as_ref()
132                .or(local_fs.as_ref())
133                .map(Arc::clone)
134                .unwrap_or_else(|| {
135                    unreachable!("local_fs is set when any glob lacks a filesystem")
136                });
137            let files: Vec<FileListing> = fs.glob(glob)?.try_collect().await?;
138            for file in files {
139                all_files.push((file, Arc::clone(&fs)));
140            }
141        }
142
143        if all_files.is_empty() {
144            let globs: Vec<_> = self.glob_sources.iter().map(|(g, _)| g.as_str()).collect();
145            vortex_bail!("No files matched the glob pattern(s): {:?}", globs);
146        }
147
148        let file_count = all_files.len();
149        let globs: Vec<_> = self.glob_sources.iter().map(|(g, _)| g.as_str()).collect();
150        debug!(file_count, glob = ?globs, "discovered files");
151
152        // Open first file eagerly for dtype.
153        let (first_file_listing, first_fs) = &all_files[0];
154        let open_fn = self.open_options_fn.as_ref();
155        let first_file = open_file(first_fs, first_file_listing, &self.session, open_fn).await?;
156        let first_reader = layout_reader_with_stats(&first_file)?;
157
158        let factories: Vec<Arc<dyn LayoutReaderFactory>> = all_files[1..]
159            .iter()
160            .map(|(file, fs)| {
161                Arc::new(VortexFileReaderFactory {
162                    fs: Arc::clone(fs),
163                    file: file.clone(),
164                    session: self.session.clone(),
165                    open_options_fn: Arc::clone(&self.open_options_fn),
166                }) as Arc<dyn LayoutReaderFactory>
167            })
168            .collect();
169
170        let inner = MultiLayoutDataSource::new_with_first(first_reader, factories, &self.session);
171
172        debug!(file_count, dtype = %inner.dtype(), "built MultiFileDataSource");
173
174        Ok(inner)
175    }
176}
177
178/// Creates a local filesystem backed by `object_store::local::LocalFileSystem`.
179// TODO(ngates): create a native file system without an object_store dependency.
180//  Turns out it's not a trivial change because we have always used object_store with its own
181//  coalescing and concurrency configs, so we need to re-tune for local disk.
182#[cfg(feature = "object_store")]
183fn create_local_filesystem(session: &VortexSession) -> VortexResult<FileSystemRef> {
184    use vortex_io::object_store::ObjectStoreFileSystem;
185    use vortex_io::session::RuntimeSessionExt;
186
187    let store = Arc::new(object_store::local::LocalFileSystem::default());
188    let fs: FileSystemRef = Arc::new(ObjectStoreFileSystem::new(store, session.handle()));
189    Ok(fs)
190}
191
192#[cfg(not(feature = "object_store"))]
193fn create_local_filesystem(_session: &VortexSession) -> VortexResult<FileSystemRef> {
194    vortex_bail!(
195        "The 'object_store' feature is required for automatic local filesystem creation. \
196             Either enable the feature or provide a filesystem via .with_filesystem()."
197    );
198}
199
200/// Open a single Vortex file, checking the session's footer cache.
201async fn open_file(
202    fs: &FileSystemRef,
203    file: &FileListing,
204    session: &VortexSession,
205    open_options_fn: &(dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync),
206) -> VortexResult<crate::VortexFile> {
207    debug!(path = %file.path, "opening vortex file");
208
209    // Open the reader first so we can use its URI as the cache key.
210    // The URI includes the full path (with any filesystem prefix), making it unique
211    // even when different PrefixFileSystem instances strip paths to the same relative name.
212    let source = fs.open_read(&file.path).await?;
213    let cache_key = source
214        .uri()
215        .map(|u| u.to_string())
216        .unwrap_or_else(|| file.path.clone());
217
218    // Build open options. The DashMap Ref from multi_file() must not live across an await,
219    // so we scope the cache lookup in a block.
220    let options = {
221        let mut options = open_options_fn(session.open_options());
222        if let Some(size) = file.size {
223            options = options.with_file_size(size);
224        }
225        if let Some(footer) = session.multi_file().get_footer(&cache_key) {
226            options = options.with_footer(footer);
227        }
228        options
229    };
230
231    let vortex_file = options.open(source).await?;
232
233    // Store footer in cache (scoped to avoid holding the Ref across subsequent code).
234    session
235        .multi_file()
236        .put_footer(&cache_key, vortex_file.footer().clone());
237    Ok(vortex_file)
238}
239
240/// Creates a layout reader from a VortexFile, wrapping with `FileStatsLayoutReader` when
241/// file-level statistics are available.
242fn layout_reader_with_stats(file: &crate::VortexFile) -> VortexResult<LayoutReaderRef> {
243    let mut reader = file.layout_reader()?;
244    if let Some(stats) = file.file_stats().cloned() {
245        reader = Arc::new(FileStatsLayoutReader::new(
246            reader,
247            stats,
248            file.session.clone(),
249        ));
250    }
251    Ok(reader)
252}
253
254/// A [`LayoutReaderFactory`] that lazily opens a single Vortex file and returns its layout reader.
255struct VortexFileReaderFactory {
256    fs: FileSystemRef,
257    file: FileListing,
258    session: VortexSession,
259    open_options_fn: Arc<dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync>,
260}
261
262#[async_trait]
263impl LayoutReaderFactory for VortexFileReaderFactory {
264    async fn open(&self) -> VortexResult<Option<LayoutReaderRef>> {
265        let file = open_file(
266            &self.fs,
267            &self.file,
268            &self.session,
269            self.open_options_fn.as_ref(),
270        )
271        .await?;
272        Ok(Some(layout_reader_with_stats(&file)?))
273    }
274}
vortex_file/multi/mod.rs

vortex_file/multi/
mod.rs
