fixes

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

java/vortex-jni/src/test/java/dev/vortex/api/TestMinimal.java

@@ -11,6 +11,7 @@
 import dev.vortex.jni.NativeLoader;
 import java.io.IOException;
 import java.math.BigDecimal;
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -132,6 +133,7 @@ public void testFullScan() throws Exception {
         DataSource ds = DataSource.open(session, writePath);
 
         assertEquals(new DataSource.RowCount.Exact(10L), ds.rowCount());
+        assertEquals(new DataSource.ByteSize.Exact(Files.size(tempDir.resolve("minimal.vortex"))), ds.byteSize());
 
         var schema = ds.arrowSchema(allocator);
         assertEquals(

java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexBatchExec.java

@@ -4,6 +4,7 @@
 package dev.vortex.spark.read;
 
 import com.google.common.collect.ImmutableMap;
+import dev.vortex.api.Session;
 import dev.vortex.jni.NativeFiles;
 import dev.vortex.spark.VortexFilePartition;
 import dev.vortex.spark.VortexSparkSession;
@@ -76,14 +77,19 @@ public PartitionReaderFactory createReaderFactory() {
     }
 
     private List<String> resolvePaths() {
-        var session = VortexSparkSession.get(formatOptions);
+        return resolveVortexPaths(VortexSparkSession.get(formatOptions), paths, formatOptions);
+    }
+
+    /**
+     * Expands directory-like entries to concrete {@code .vortex} files; entries that already name a {@code .vortex}
+     * file are kept as-is. Shared with {@link VortexScan#estimateStatistics()} so planning and execution resolve paths
+     * identically.
+     */
+    static List<String> resolveVortexPaths(Session session, List<String> paths, Map<String, String> formatOptions) {
         return paths.stream()
-                .flatMap(path -> {
-                    if (path.endsWith(".vortex")) {
-                        return Stream.of(path);
-                    }
-                    return NativeFiles.listFiles(session, path, formatOptions).stream();
-                })
+                .flatMap(path -> path.endsWith(".vortex")
+                        ? Stream.of(path)
+                        : NativeFiles.listFiles(session, path, formatOptions).stream())
                 .collect(Collectors.toList());
     }
 

java/vortex-spark/src/main/java/dev/vortex/spark/read/VortexScan.java

@@ -5,15 +5,11 @@
 
 import dev.vortex.api.DataSource;
 import dev.vortex.api.Session;
-import dev.vortex.jni.NativeFiles;
 import dev.vortex.spark.VortexSparkSession;
 import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.OptionalLong;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.apache.spark.sql.connector.catalog.CatalogV2Util;
 import org.apache.spark.sql.connector.catalog.Column;
 import org.apache.spark.sql.connector.expressions.NamedReference;
@@ -140,14 +136,7 @@ public Statistics estimateStatistics() {
 
     private Statistics computeStatistics() {
         Session session = VortexSparkSession.get(formatOptions);
-        // Expand directory paths to concrete files the way VortexBatchExec does, so we use the
-        // same per-path resolution end-to-end.
-        List<String> resolvedPaths = paths.stream()
-                .flatMap(path -> path.endsWith(".vortex")
-                        ? Stream.of(path)
-                        : NativeFiles.listFiles(session, path, formatOptions).stream())
-                .collect(Collectors.toList());
-
+        List<String> resolvedPaths = VortexBatchExec.resolveVortexPaths(session, paths, formatOptions);
         if (resolvedPaths.isEmpty()) {
             return new VortexStatistics(OptionalLong.empty(), OptionalLong.empty());
         }
@@ -181,7 +170,7 @@ private record VortexStatistics(OptionalLong numRows, OptionalLong sizeInBytes)
 
         @Override
         public Map<NamedReference, ColumnStatistics> columnStats() {
-            return new HashMap<>();
+            return Map.of();
         }
     }
 }

java/vortex-spark/src/test/java/dev/vortex/spark/VortexDataSourceStatsTest.java

@@ -98,17 +98,26 @@ public void testEstimateStatisticsAcrossMultipleFiles() throws IOException {
     public void testEstimateStatisticsReportsSizeInBytes() throws IOException {
         Path outputPath = writeRows(120, "with_size", 3);
 
-        long expectedTotalBytes = totalVortexFileBytes(outputPath);
-        assertTrue(expectedTotalBytes > 0, "Test setup should produce at least one non-empty .vortex file");
+        long fileBytes = totalVortexFileBytes(outputPath);
+        assertTrue(fileBytes > 0, "Test setup should produce at least one non-empty .vortex file");
 
         VortexScan scan = buildScan(outputPath);
         Statistics stats = scan.estimateStatistics();
 
         assertTrue(
                 stats.sizeInBytes().isPresent(),
                 "VortexScan should surface a sizeInBytes when the filesystem listing reports file sizes");
+        // Mirror the scan's Spark-convention scaling (factor 1.0, unpruned schema), which divides and
+        // re-multiplies by the schema default size in double arithmetic before truncating; asserting
+        // against the raw byte sum would be sensitive to the floating-point round trip.
+        StructType schema = spark.read()
+                .format("vortex")
+                .option("path", outputPath.toUri().toString())
+                .load()
+                .schema();
+        long expectedSize = (long) (1.0 * fileBytes / schema.defaultSize() * schema.defaultSize());
         assertEquals(
-                expectedTotalBytes,
+                expectedSize,
                 stats.sizeInBytes().getAsLong(),
                 "sizeInBytes should equal the sum of on-storage .vortex file sizes");
     }

vortex-file/src/multi/mod.rs

@@ -8,9 +8,12 @@ mod session;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use futures::StreamExt;
 use futures::TryStreamExt;
+use futures::stream;
 use session::MultiFileSessionExt;
 use tracing::debug;
+use vortex_error::VortexError;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_io::filesystem::FileListing;
@@ -60,19 +63,20 @@ pub struct MultiFileDataSource {
     /// List of (glob, optional filesystem) pairs to resolve.
     /// When the filesystem is None, a local filesystem will be created in build().
     glob_sources: Vec<(String, Option<FileSystemRef>)>,
-    /// Pre-resolved file listings that skip glob expansion. The caller is responsible for
-    /// supplying the [`FileListing::size`] when stats reporting matters.
-    listing_sources: Vec<(FileListing, FileSystemRef)>,
     open_options_fn: Arc<dyn Fn(VortexOpenOptions) -> VortexOpenOptions + Send + Sync>,
 }
 
+/// In-flight glob resolutions in [`MultiFileDataSource::build`]. Callers like the JNI data
+/// source add one exact path per glob source, where each resolution is a single remote
+/// metadata lookup; resolving them concurrently avoids one round trip of latency per file.
+const GLOB_RESOLUTION_CONCURRENCY: usize = 16;
+
 impl MultiFileDataSource {
     /// Create a new [`MultiFileDataSource`] builder.
     pub fn new(session: VortexSession) -> Self {
         Self {
             session,
             glob_sources: Vec::new(),
-            listing_sources: Vec::new(),
             open_options_fn: Arc::new(|opts| opts),
         }
     }
@@ -98,23 +102,6 @@ impl MultiFileDataSource {
         self
     }
 
-    /// Add a pre-resolved file listing.
-    ///
-    /// Use this when the caller already knows the exact file path and (optionally) its size,
-    /// avoiding the glob expansion done by [`Self::with_glob`]. Supplying
-    /// [`FileListing::size`] is required for [`DataSource::byte_size`] to surface a contribution
-    /// from this file; otherwise the source size remains unknown for this file and the
-    /// data-source-level total is extrapolated from the files that do report a size.
-    pub fn with_listing(mut self, listing: FileListing, fs: FileSystemRef) -> Self {
-        let FileListing { path, size } = listing;
-        let listing = FileListing {
-            path: path.trim_start_matches('/').to_string(),
-            size,
-        };
-        self.listing_sources.push((listing, fs));
-        self
-    }
-
     /// Customize [`VortexOpenOptions`] applied to each file.
     ///
     /// Use this to configure segment caches, metrics registries, or other per-file options.
@@ -131,10 +118,8 @@ impl MultiFileDataSource {
     /// Discovers files via glob, opens the first file eagerly to determine the schema,
     /// and creates lazy factories for the remaining files.
     pub async fn build(self) -> VortexResult<MultiLayoutDataSource> {
-        if self.glob_sources.is_empty() && self.listing_sources.is_empty() {
-            vortex_bail!(
-                "MultiFileDataSource requires at least one glob pattern or pre-resolved listing"
-            );
+        if self.glob_sources.is_empty() {
+            vortex_bail!("MultiFileDataSource requires at least one glob pattern");
         }
 
         // Create local filesystem lazily if needed (only if any glob lacks a filesystem).
@@ -145,32 +130,39 @@ impl MultiFileDataSource {
             .then(|| create_local_filesystem(&self.session))
             .transpose()?;
 
-        // Collect files from all glob sources.
-        let mut all_files: Vec<(FileListing, FileSystemRef)> = Vec::new();
-        for (glob, maybe_fs) in &self.glob_sources {
-            // Use the provided filesystem, or fall back to the local filesystem.
-            // We know local_fs is Some when maybe_fs is None (by construction above).
-            let fs = maybe_fs
-                .as_ref()
-                .or(local_fs.as_ref())
-                .map(Arc::clone)
-                .unwrap_or_else(|| {
-                    unreachable!("local_fs is set when any glob lacks a filesystem")
-                });
-            let files: Vec<FileListing> = fs.glob(glob)?.try_collect().await?;
-            for file in files {
-                all_files.push((file, Arc::clone(&fs)));
-            }
-        }
-        all_files.extend(self.listing_sources);
+        let globs: Vec<String> = self.glob_sources.iter().map(|(g, _)| g.clone()).collect();
+
+        // Resolve glob sources concurrently while preserving their order, since the order
+        // determines partition indices and which file is opened eagerly for the schema.
+        let resolved: Vec<Vec<(FileListing, FileSystemRef)>> =
+            stream::iter(self.glob_sources.into_iter().map(|(glob, maybe_fs)| {
+                // Use the provided filesystem, or fall back to the local filesystem.
+                // We know local_fs is Some when maybe_fs is None (by construction above).
+                let fs = maybe_fs
+                    .or_else(|| local_fs.as_ref().map(Arc::clone))
+                    .unwrap_or_else(|| {
+                        unreachable!("local_fs is set when any glob lacks a filesystem")
+                    });
+                async move {
+                    let files: Vec<FileListing> = fs.glob(&glob)?.try_collect().await?;
+                    Ok::<_, VortexError>(
+                        files
+                            .into_iter()
+                            .map(|file| (file, Arc::clone(&fs)))
+                            .collect(),
+                    )
+                }
+            }))
+            .buffered(GLOB_RESOLUTION_CONCURRENCY)
+            .try_collect()
+            .await?;
+        let all_files: Vec<(FileListing, FileSystemRef)> = resolved.into_iter().flatten().collect();
 
         if all_files.is_empty() {
-            let globs: Vec<_> = self.glob_sources.iter().map(|(g, _)| g.as_str()).collect();
             vortex_bail!("No files matched the glob pattern(s): {:?}", globs);
         }
 
         let file_count = all_files.len();
-        let globs: Vec<_> = self.glob_sources.iter().map(|(g, _)| g.as_str()).collect();
         debug!(file_count, glob = ?globs, "discovered files");
 
         // Open first file eagerly for dtype.

vortex-jni/src/data_source.rs

@@ -16,8 +16,6 @@ use std::path::PathBuf;
 use std::path::absolute;
 use std::sync::Arc;
 
-use futures::StreamExt;
-use futures::stream;
 use jni::EnvUnowned;
 use jni::objects::JClass;
 use jni::objects::JLongArray;
@@ -30,7 +28,6 @@ use vortex::error::VortexResult;
 use vortex::error::vortex_err;
 use vortex::expr::stats::Precision;
 use vortex::file::multi::MultiFileDataSource;
-use vortex::io::filesystem::FileListing;
 use vortex::io::filesystem::FileSystemRef;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::session::RuntimeSessionExt;
@@ -44,10 +41,6 @@ use crate::file::extract_properties;
 use crate::object_store::object_store_fs;
 use crate::session::session_ref;
 
-/// In-flight size lookups while resolving exact paths to file listings. Balances HEAD
-/// throughput on remote stores against connection overhead.
-const SIZE_LOOKUP_CONCURRENCY: usize = 16;
-
 /// Wraps an `Arc<dyn DataSource>` behind a single pointer.
 pub(crate) struct NativeDataSource {
     inner: DataSourceRef,
@@ -110,49 +103,14 @@ pub extern "system" fn Java_dev_vortex_jni_NativeDataSource_open(
             }
         }
 
-        // Split inputs into glob patterns (which fs.glob() expands via list(), capturing sizes
-        // automatically) and exact paths (which are resolved one-by-one with a HEAD-style size
-        // lookup so the data source can report total bytes for Spark-style stats).
-        let mut glob_inputs: Vec<(String, FileSystemRef)> = Vec::new();
-        let mut exact_inputs: Vec<(String, FileSystemRef)> = Vec::new();
+        let mut builder = MultiFileDataSource::new(session.clone());
         for glob_url in &glob_urls {
             let base = base_url(glob_url);
             let fs = fs_cache
                 .get(&base)
                 .cloned()
                 .unwrap_or_else(|| unreachable!("fs cached for every base url"));
-            let path = glob_url.path().to_string();
-            if path.contains(['*', '?', '[']) {
-                glob_inputs.push((path, fs));
-            } else {
-                exact_inputs.push((path, fs));
-            }
-        }
-
-        let resolved_listings: Vec<(FileListing, FileSystemRef)> = if exact_inputs.is_empty() {
-            Vec::new()
-        } else {
-            RUNTIME.block_on(async {
-                stream::iter(exact_inputs)
-                    .map(|(path, fs)| async move {
-                        let size = match fs.open_read(&path).await {
-                            Ok(source) => source.size().await.ok(),
-                            Err(_) => None,
-                        };
-                        (FileListing { path, size }, fs)
-                    })
-                    .buffer_unordered(SIZE_LOOKUP_CONCURRENCY)
-                    .collect::<Vec<_>>()
-                    .await
-            })
-        };
-
-        let mut builder = MultiFileDataSource::new(session.clone());
-        for (glob, fs) in glob_inputs {
-            builder = builder.with_glob(glob, Some(fs));
-        }
-        for (listing, fs) in resolved_listings {
-            builder = builder.with_listing(listing, fs);
+            builder = builder.with_glob(glob_url.path(), Some(fs));
         }
 
         let inner = RUNTIME

vortex-layout/src/scan/multi.rs

@@ -530,3 +530,43 @@ impl Partition for MultiLayoutPartition {
         )))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+    use vortex_array::dtype::Nullability;
+
+    use super::*;
+    use crate::scan::test::new_session;
+
+    struct NeverOpened;
+
+    #[async_trait]
+    impl LayoutReaderFactory for NeverOpened {
+        async fn open(&self) -> VortexResult<Option<LayoutReaderRef>> {
+            unreachable!("byte_size must not open readers")
+        }
+    }
+
+    fn deferred_source(byte_sizes: Vec<Option<u64>>) -> MultiLayoutDataSource {
+        let factories: Vec<Arc<dyn LayoutReaderFactory>> = byte_sizes
+            .iter()
+            .map(|_| Arc::new(NeverOpened) as _)
+            .collect();
+        MultiLayoutDataSource::new_deferred(
+            DType::Bool(Nullability::NonNullable),
+            factories,
+            byte_sizes,
+            &new_session(),
+        )
+    }
+
+    #[rstest]
+    #[case::all_known(vec![Some(10), Some(20), Some(30)], Precision::exact(60u64))]
+    #[case::some_known_extrapolates(vec![Some(10), None, Some(30)], Precision::inexact(60u64))]
+    #[case::none_known(vec![None, None], Precision::Absent)]
+    #[case::no_children(vec![], Precision::exact(0u64))]
+    fn byte_size_precision(#[case] sizes: Vec<Option<u64>>, #[case] expected: Precision<u64>) {
+        assert_eq!(deferred_source(sizes).byte_size(), expected);
+    }
+}