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

vortex_file/multi/
mod.rs
