Make footer and segment cache sizes configurable in Datafusion (#4491)

AdamGS · web-flow · commit d135f4a1993b · 2025-09-03T12:28:44.000+01:00
Aside from allowing configuring these two values now, this is a step
towards a more flexible integration with DataFusion.

Signed-off-by: Adam Gutglick &lt;adam@spiraldb.com&gt;
diff --git a/bench-vortex/src/engines/df/mod.rs b/bench-vortex/src/engines/df/mod.rs
@@ -66,7 +66,7 @@ pub fn get_session_context(disable_datafusion_cache: bool) -> SessionContext {
         .build_arc()
         .expect("could not build runtime environment");
 
-    let factory = VortexFormatFactory::default();
+    let factory = VortexFormatFactory::new();
 
     let mut session_state_builder = SessionStateBuilder::new()
         .with_config(SessionConfig::default())
diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs
@@ -8,11 +8,12 @@ use std::sync::Arc;
 use arrow_schema::{Schema, SchemaRef};
 use async_trait::async_trait;
 use datafusion_catalog::Session;
+use datafusion_common::config::ConfigField;
 use datafusion_common::parsers::CompressionTypeVariant;
 use datafusion_common::stats::Precision;
 use datafusion_common::{
-    ColumnStatistics, DataFusionError, GetExt, Result as DFResult, Statistics,
-    config_datafusion_err, not_impl_err,
+    ColumnStatistics, DataFusionError, GetExt, Result as DFResult, Statistics, config_namespace,
+    not_impl_err,
 };
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_datasource::file::FileSource;
@@ -48,7 +49,7 @@ use crate::convert::TryToDataFusion;
 pub struct VortexFormat {
     session: Arc<VortexSession>,
     file_cache: VortexFileCache,
-    opts: VortexFormatOptions,
+    opts: VortexOptions,
 }
 
 impl Debug for VortexFormat {
@@ -59,28 +60,27 @@ impl Debug for VortexFormat {
     }
 }
 
-/// Options to configure the [`VortexFormat`].
-#[derive(Debug)]
-pub struct VortexFormatOptions {
-    /// The size of the in-memory [`vortex::file::Footer`] cache.
-    pub footer_cache_size_mb: usize,
-    /// The size of the in-memory segment cache.
-    pub segment_cache_size_mb: usize,
-}
-
-impl Default for VortexFormatOptions {
-    fn default() -> Self {
-        Self {
-            footer_cache_size_mb: 64,
-            segment_cache_size_mb: 0,
-        }
+config_namespace! {
+    /// Options to configure the [`VortexFormat`].
+    ///
+    /// Can be set through a DataFusion [`SessionConfig`].
+    ///
+    /// [`SessionConfig`]: https://docs.rs/datafusion/latest/datafusion/prelude/struct.SessionConfig.html
+    pub struct VortexOptions {
+        /// The size of the in-memory [`vortex::file::Footer`] cache.
+        pub footer_cache_size_mb: usize, default = 64
+        /// The size of the in-memory segment cache.
+        pub segment_cache_size_mb: usize, default = 0
     }
 }
 
+impl Eq for VortexOptions {}
+
 /// Minimal factory to create [`VortexFormat`] instances.
-#[derive(Default, Debug)]
+#[derive(Debug)]
 pub struct VortexFormatFactory {
     session: Arc<VortexSession>,
+    options: Option<VortexOptions>,
 }
 
 impl GetExt for VortexFormatFactory {
@@ -89,20 +89,60 @@ impl GetExt for VortexFormatFactory {
     }
 }
 
+impl VortexFormatFactory {
+    /// Creates a new instance with a default [`VortexSession`] and default options.
+    #[allow(clippy::new_without_default)] // FormatFactory defines `default` method, so having `Default` implementation is confusing.
+    pub fn new() -> Self {
+        Self {
+            session: Arc::new(VortexSession::default()),
+            options: None,
+        }
+    }
+
+    /// Creates a new instance with customized session and default options for all [`VortexFormat`] instances created from this factory.
+    ///
+    /// The options can be overridden by table-level configuration pass in [`FileFormatFactory::create`].
+    pub fn new_with_options(session: Arc<VortexSession>, options: VortexOptions) -> Self {
+        Self {
+            session,
+            options: Some(options),
+        }
+    }
+
+    /// Override the default options for this factory.
+    ///
+    /// For example:
+    /// ```rust
+    /// use vortex_datafusion::{VortexFormatFactory, VortexOptions};
+    ///
+    /// let factory = VortexFormatFactory::new().with_options(VortexOptions::default());
+    /// ```
+    pub fn with_options(mut self, options: VortexOptions) -> Self {
+        self.options = Some(options);
+        self
+    }
+}
+
 impl FileFormatFactory for VortexFormatFactory {
     #[allow(clippy::disallowed_types)]
     fn create(
         &self,
         _state: &dyn Session,
         format_options: &std::collections::HashMap<String, String>,
     ) -> DFResult<Arc<dyn FileFormat>> {
-        if !format_options.is_empty() {
-            return Err(config_datafusion_err!(
-                "Vortex tables don't support any options"
-            ));
+        let mut opts = self.options.clone().unwrap_or_default();
+        for (key, value) in format_options {
+            if let Some(key) = key.strip_prefix("format.") {
+                opts.set(key, value)?;
+            } else {
+                tracing::trace!("Ignoring options '{key}'");
+            }
         }
 
-        Ok(Arc::new(VortexFormat::new(self.session.clone())))
+        Ok(Arc::new(VortexFormat::new_with_options(
+            self.session.clone(),
+            opts,
+        )))
     }
 
     fn default(&self) -> Arc<dyn FileFormat> {
@@ -121,9 +161,13 @@ impl Default for VortexFormat {
 }
 
 impl VortexFormat {
-    /// Create a new instance of the [`VortexFormat`].
+    /// Create a new instance with default options.
     pub fn new(session: Arc<VortexSession>) -> Self {
-        let opts = VortexFormatOptions::default();
+        Self::new_with_options(session, VortexOptions::default())
+    }
+
+    /// Creates a new instance with configured by a [`VortexOptions`].
+    pub fn new_with_options(session: Arc<VortexSession>, opts: VortexOptions) -> Self {
         Self {
             session: session.clone(),
             file_cache: VortexFileCache::new(
@@ -136,7 +180,7 @@ impl VortexFormat {
     }
 
     /// Return the format specific configuration
-    pub fn options(&self) -> &VortexFormatOptions {
+    pub fn options(&self) -> &VortexOptions {
         &self.opts
     }
 }
@@ -389,7 +433,7 @@ mod tests {
     async fn create_table() {
         let dir = TempDir::new().unwrap();
 
-        let factory: VortexFormatFactory = Default::default();
+        let factory: VortexFormatFactory = VortexFormatFactory::new();
         let mut session_state_builder = SessionStateBuilder::new().with_default_features();
         register_vortex_format_factory(factory, &mut session_state_builder);
         let session = SessionContext::new_with_state(session_state_builder.build());
@@ -408,11 +452,10 @@ mod tests {
     }
 
     #[tokio::test]
-    #[should_panic]
-    async fn fail_table_config() {
+    async fn configure_format_source() {
         let dir = TempDir::new().unwrap();
 
-        let factory: VortexFormatFactory = Default::default();
+        let factory = VortexFormatFactory::new();
         let mut session_state_builder = SessionStateBuilder::new().with_default_features();
         register_vortex_format_factory(factory, &mut session_state_builder);
         let session = SessionContext::new_with_state(session_state_builder.build());
@@ -422,7 +465,7 @@ mod tests {
                 "CREATE EXTERNAL TABLE my_tbl \
                 (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
                 STORED AS vortex LOCATION '{}' \
-                OPTIONS( some_key 'value' );",
+                OPTIONS( segment_cache_size_mb '5' );",
                 dir.path().to_str().unwrap()
             ))
             .await
diff --git a/vortex-datafusion/src/persistent/mod.rs b/vortex-datafusion/src/persistent/mod.rs
@@ -9,7 +9,7 @@ mod opener;
 mod sink;
 mod source;
 
-pub use format::{VortexFormat, VortexFormatFactory, VortexFormatOptions};
+pub use format::{VortexFormat, VortexFormatFactory, VortexOptions};
 pub use source::VortexSource;
 
 #[cfg(test)]
diff --git a/vortex-datafusion/src/persistent/sink.rs b/vortex-datafusion/src/persistent/sink.rs
@@ -172,7 +172,8 @@ mod tests {
     async fn test_insert_into() {
         let dir = TempDir::new().unwrap();
 
-        let factory = VortexFormatFactory::default();
+        let factory = VortexFormatFactory::new();
+
         let mut session_state_builder = SessionStateBuilder::new().with_default_features();
         register_vortex_format_factory(factory, &mut session_state_builder);
         let session = SessionContext::new_with_state(session_state_builder.build());
@@ -181,7 +182,7 @@ mod tests {
             .sql(&format!(
                 "CREATE EXTERNAL TABLE my_tbl \
                     (c1 VARCHAR NOT NULL, c2 INT NOT NULL) \
-                STORED AS vortex 
+                STORED AS vortex \
                 LOCATION '{}/';",
                 dir.path().to_str().unwrap()
             ))
@@ -255,7 +256,8 @@ mod tests {
 
         let dir = TempDir::new()?;
 
-        let factory = VortexFormatFactory::default();
+        let factory = VortexFormatFactory::new();
+
         let mut session_state_builder = SessionStateBuilder::new().with_default_features();
         register_vortex_format_factory(factory, &mut session_state_builder);
         let session = SessionContext::new_with_state(session_state_builder.build());
@@ -268,7 +270,7 @@ mod tests {
         let logical_plan = LogicalPlanBuilder::copy_to(
             data.logical_plan().clone(),
             dir.path().to_str().unwrap().to_string(),
-            format_as_file_type(Arc::new(VortexFormatFactory::default())),
+            format_as_file_type(Arc::new(VortexFormatFactory::new())),
             Default::default(),
             vec![],
         )?
@@ -285,7 +287,7 @@ mod tests {
             .sql(&format!(
                 "CREATE EXTERNAL TABLE written_data \
                     (a TINYINT NOT NULL) \
-                STORED AS vortex 
+                STORED AS vortex \
                 LOCATION '{}/';",
                 dir.path().to_str().unwrap()
             ))
