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

vortex_file/multi/
mod.rs
