Make opening a Vortex file with a cached footer cheaper (#10358)

Every DataFusion partition reopens each file it scans, even when the
footer is cached. Some of the work done on each of those opens grows
with the file's segment count, or with the square of its column count.
On index-driven scans the open dominates the query, so this overhead showed
up directly in query latency.

Author: joseph-isaacs

vortex-datafusion/src/convert/schema.rs

@@ -33,12 +33,20 @@ pub fn calculate_physical_schema(
         ));
     };
 
+    let reference_fields = reference_logical_schema.fields();
     let fields: Vec<Field> = struct_dtype
         .names()
         .iter()
         .zip(struct_dtype.fields())
-        .map(|(name, field_dtype)| {
-            let logical_field = reference_logical_schema.field_with_name(name.as_ref()).ok();
+        .enumerate()
+        .map(|(idx, (name, field_dtype))| {
+            // Files usually share the table's column order, so try the field at the same position
+            // before searching by name, which would make wide schemas quadratic.
+            let logical_field = reference_fields
+                .get(idx)
+                .filter(|field| field.name() == name.as_ref())
+                .map(|field| field.as_ref())
+                .or_else(|| reference_logical_schema.field_with_name(name.as_ref()).ok());
             match logical_field {
                 Some(logical_field) => {
                     let arrow_type = calculate_physical_field_type(

vortex-datafusion/src/persistent/opener.rs

@@ -8,6 +8,7 @@ use std::sync::Weak;
 use arrow_array::RecordBatchOptions;
 use arrow_schema::Field;
 use arrow_schema::Schema;
+use arrow_schema::SchemaRef;
 use datafusion_common::DataFusionError;
 use datafusion_common::Result as DFResult;
 use datafusion_common::ScalarValue;
@@ -97,6 +98,8 @@ pub(crate) struct VortexOpener {
     pub layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
     /// Shared full-file natural splits keyed by file path.
     pub natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
+    /// Shared physical Arrow schemas of each file, keyed by file path.
+    pub file_schemas: Arc<DashMap<Path, SchemaRef>>,
     /// Whether the query has output ordering specified
     pub has_output_ordering: bool,
 
@@ -138,6 +141,7 @@ impl FileOpener for VortexOpener {
         let limit = self.limit;
         let layout_readers = Arc::clone(&self.layout_readers);
         let natural_splits = Arc::clone(&self.natural_splits);
+        let file_schemas = Arc::clone(&self.file_schemas);
         let has_output_ordering = self.has_output_ordering;
         let scan_concurrency = self.scan_concurrency;
 
@@ -238,12 +242,20 @@ impl FileOpener for VortexOpener {
             }
 
             // This is the expected arrow types of the actual columns in the file, which might have different types
-            // from the unified logical schema or miss
-            let this_file_schema = Arc::new(calculate_physical_schema(
-                vxf.dtype(),
-                &unified_file_schema,
-                &session.arrow(),
-            )?);
+            // from the unified logical schema or miss. Every partition opening this file computes
+            // the same schema, so it is shared across them.
+            let this_file_schema = match file_schemas.get(&file.object_meta.location) {
+                Some(schema) => Arc::clone(schema.value()),
+                None => {
+                    let schema = Arc::new(calculate_physical_schema(
+                        vxf.dtype(),
+                        &unified_file_schema,
+                        &session.arrow(),
+                    )?);
+                    file_schemas.insert(file.object_meta.location.clone(), Arc::clone(&schema));
+                    schema
+                }
+            };
 
             let expr_adapter = expr_adapter_factory.create(
                 Arc::clone(&unified_file_schema),
@@ -865,6 +877,7 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             natural_splits: Default::default(),
+            file_schemas: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1148,6 +1161,7 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             natural_splits: Default::default(),
+            file_schemas: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1235,6 +1249,7 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             natural_splits: Default::default(),
+            file_schemas: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1389,6 +1404,7 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             natural_splits: Default::default(),
+            file_schemas: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1449,6 +1465,7 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             natural_splits: Default::default(),
+            file_schemas: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,
@@ -1656,6 +1673,7 @@ mod tests {
             df_metrics: ExecutionPlanMetricsSet::new(),
             layout_readers: Default::default(),
             natural_splits: Default::default(),
+            file_schemas: Default::default(),
             has_output_ordering: false,
             expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
             file_metadata_cache: None,

vortex-datafusion/src/persistent/source.rs

@@ -5,6 +5,7 @@ use std::fmt::Formatter;
 use std::sync::Arc;
 use std::sync::Weak;
 
+use arrow_schema::SchemaRef;
 use datafusion_common::Result as DFResult;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::TreeNodeRecursion;
@@ -198,6 +199,8 @@ pub struct VortexSource {
     layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
     /// Shared full-file natural splits keyed by path.
     natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
+    /// Shared physical Arrow schemas of each file, keyed by path.
+    file_schemas: Arc<DashMap<Path, SchemaRef>>,
     expression_convertor: Arc<dyn ExpressionConvertor>,
     pub(crate) vortex_reader_factory: Option<Arc<dyn VortexReaderFactory>>,
     pub(crate) ordered: bool,
@@ -231,6 +234,7 @@ impl VortexSource {
             df_metrics: Default::default(),
             layout_readers: Arc::new(DashMap::default()),
             natural_splits: Arc::new(DashMap::default()),
+            file_schemas: Arc::new(DashMap::default()),
             expression_convertor,
             vortex_reader_factory: None,
             vx_metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
@@ -355,6 +359,7 @@ impl VortexSource {
             df_metrics: self.df_metrics.clone(),
             layout_readers: Arc::clone(&self.layout_readers),
             natural_splits: Arc::clone(&self.natural_splits),
+            file_schemas: Arc::clone(&self.file_schemas),
             has_output_ordering: !base_config.output_ordering.is_empty() || self.ordered,
             expression_convertor: Arc::clone(&self.expression_convertor),
             file_metadata_cache: self.file_metadata_cache.clone(),

vortex-file/src/footer/mod.rs

@@ -61,6 +61,14 @@ pub struct Footer {
     segments: Arc<[SegmentSpec]>,
     statistics: Option<FileStatistics>,
     metadata: Arc<[(String, SegmentSpec)]>,
+    // `segments` followed by the `metadata` segments, which is the segment table a file is opened
+    // with. Precomputed so opening a file with a cached footer does not copy every segment spec.
+    segments_with_metadata: Arc<[SegmentSpec]>,
+    // The largest end offset of any segment in `segments`, or `None` if one overflows `u64`.
+    // Precomputed so validating a cached footer against a file size does not walk every segment.
+    // Writers never overlap segments, so in a valid file this is the last segment's end, but only
+    // their order is checked when reading, and this guards against malformed footers.
+    segments_end: Option<u64>,
     // The specific arrays used within the file, in the order they were registered.
     array_read_ctx: ReadContext,
     // The approximate size of the footer in bytes, used for caching and memory management.
@@ -76,6 +84,8 @@ impl Footer {
     ) -> Self {
         Self {
             root_layout,
+            segments_with_metadata: Arc::clone(&segments),
+            segments_end: segments_end(&segments),
             segments,
             statistics,
             metadata: Arc::from([]),
@@ -148,6 +158,8 @@ impl Footer {
 
         Ok(Self {
             root_layout,
+            segments_with_metadata: segments_with_metadata(&segments, &metadata),
+            segments_end: segments_end(&segments),
             segments,
             statistics,
             metadata,
@@ -186,16 +198,13 @@ impl Footer {
     }
 
     pub(crate) fn with_metadata_segments(mut self, metadata: Arc<[(String, SegmentSpec)]>) -> Self {
+        self.segments_with_metadata = segments_with_metadata(&self.segments, &metadata);
         self.metadata = metadata;
         self
     }
 
     pub(crate) fn segment_specs_with_metadata(&self) -> Arc<[SegmentSpec]> {
-        self.segments
-            .iter()
-            .copied()
-            .chain(self.metadata.iter().map(|(_, segment)| *segment))
-            .collect()
+        Arc::clone(&self.segments_with_metadata)
     }
 
     /// Computes the compressed size in bytes of every field in the file, keyed by field path.
@@ -224,6 +233,10 @@ impl Footer {
 
     /// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
     pub(crate) fn validate_file_size(&self, file_size: u64) -> VortexResult<()> {
+        if self.segments_end.is_some_and(|end| end <= file_size) {
+            return Ok(());
+        }
+        // Walk the segments only to name the offending one in the error.
         validate_segments_within_file(&self.segments, file_size)
     }
 
@@ -238,6 +251,26 @@ impl Footer {
     }
 }
 
+fn segments_with_metadata(
+    segments: &Arc<[SegmentSpec]>,
+    metadata: &[(String, SegmentSpec)],
+) -> Arc<[SegmentSpec]> {
+    if metadata.is_empty() {
+        return Arc::clone(segments);
+    }
+    segments
+        .iter()
+        .copied()
+        .chain(metadata.iter().map(|(_, segment)| *segment))
+        .collect()
+}
+
+fn segments_end(segments: &[SegmentSpec]) -> Option<u64> {
+    segments.iter().try_fold(0u64, |end, segment| {
+        Some(end.max(segment.offset.checked_add(u64::from(segment.length))?))
+    })
+}
+
 /// Validate that every segment declared in the footer lies within a file of `file_size` bytes.
 ///
 /// A corrupt or malicious file can declare a segment whose offset or length extends past the end
@@ -293,4 +326,29 @@ mod tests {
         let err = validate_segments_within_file(&[segment(u64::MAX, 1)], u64::MAX).unwrap_err();
         assert!(err.to_string().contains("past the end"), "{err}");
     }
+
+    #[test]
+    fn segments_end_covers_overlapping_segments() {
+        // Valid files never overlap segments, but a malformed footer may declare an earlier
+        // segment that ends after a later one, and the end must cover it so that file size
+        // validation still rejects it.
+        assert_eq!(
+            segments_end(&[segment(0, 300), segment(100, 50)]),
+            Some(300)
+        );
+        assert_eq!(segments_end(&[]), Some(0));
+        assert_eq!(segments_end(&[segment(0, 1), segment(u64::MAX, 1)]), None);
+    }
+
+    #[test]
+    fn segments_with_metadata_appends_metadata_segments() {
+        let segments: Arc<[SegmentSpec]> = Arc::from([segment(0, 10), segment(10, 20)]);
+
+        let without = segments_with_metadata(&segments, &[]);
+        assert!(Arc::ptr_eq(&without, &segments));
+
+        let with = segments_with_metadata(&segments, &[("key".to_string(), segment(30, 5))]);
+        let spans: Vec<_> = with.iter().map(|s| (s.offset, s.length)).collect();
+        assert_eq!(spans, [(0, 10), (10, 20), (30, 5)]);
+    }
 }