feat: Support retrieving writer strategy builder from vortex session

Author: peasee

vortex-file/src/strategy.rs

@@ -23,6 +23,7 @@ const ONE_MEG: u64 = 1 << 20;
 /// Vortex provides an out-of-the-box file writer that optimizes the layout of chunks on-disk,
 /// repartitioning and compressing them to strike a balance between size on-disk,
 /// bulk decoding performance, and IOPS required to perform an indexed read.
+#[derive(Debug, Clone)]
 pub struct WriteStrategyBuilder {
     compressor: Option<Arc<dyn CompressorPlugin>>,
     row_block_size: usize,

vortex-file/src/writer.rs

@@ -45,9 +45,14 @@ pub struct VortexWriteOptions {
 pub trait WriteOptionsSessionExt: SessionExt {
     /// Create [`VortexWriteOptions`] for writing to a Vortex file.
     fn write_options(&self) -> VortexWriteOptions {
+        let maybe_write_strategy_builder = self.get_opt::<WriteStrategyBuilder>();
+        let strategy = maybe_write_strategy_builder
+            .map(|opt| opt.clone().build())
+            .unwrap_or_else(|| WriteStrategyBuilder::new().build());
+
         VortexWriteOptions {
             session: self.session(),
-            strategy: WriteStrategyBuilder::new().build(),
+            strategy,
             exclude_dtype: false,
             file_statistics: PRUNING_STATS.to_vec(),
             max_variable_length_statistics_size: 64,
@@ -431,3 +436,20 @@ impl WriteSummary {
         self.footer.row_count()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_write_options_from_session_vars() {
+        let session = VortexSession::empty();
+        let fetched_write_strategy = session.get_opt::<WriteStrategyBuilder>();
+        assert!(fetched_write_strategy.is_none());
+        drop(fetched_write_strategy);
+
+        let session = session.set(WriteStrategyBuilder::new());
+        let fetched_write_strategy = session.get_opt::<WriteStrategyBuilder>();
+        assert!(fetched_write_strategy.is_some());
+    }
+}

vortex-layout/src/layouts/compressed.rs

@@ -23,7 +23,7 @@ use crate::{LayoutRef, LayoutStrategy};
 /// meet this interface.
 ///
 /// API consumers are also free to implement this trait to provide new plugin compressors.
-pub trait CompressorPlugin: Send + Sync + 'static {
+pub trait CompressorPlugin: std::fmt::Debug + Send + Sync + 'static {
     fn compress_chunk(&self, chunk: &dyn Array) -> VortexResult<ArrayRef>;
 }
 
@@ -35,7 +35,7 @@ impl CompressorPlugin for Arc<dyn CompressorPlugin> {
 
 impl<F> CompressorPlugin for F
 where
-    F: Fn(&dyn Array) -> VortexResult<ArrayRef> + Send + Sync + 'static,
+    F: Fn(&dyn Array) -> VortexResult<ArrayRef> + Send + Sync + 'static + std::fmt::Debug,
 {
     fn compress_chunk(&self, chunk: &dyn Array) -> VortexResult<ArrayRef> {
         self(chunk)

vortex-session/src/lib.rs

@@ -46,6 +46,26 @@ impl VortexSession {
         }
         self
     }
+
+    /// Inserts a new session variable of type `V` with the supplied value.
+    ///
+    /// # Panics
+    ///
+    /// If a variable of that type already exists.
+    pub fn set<V: SessionVar>(self, val: V) -> Self {
+        match self.0.entry(TypeId::of::<V>()) {
+            Entry::Occupied(_) => {
+                vortex_panic!(
+                    "Session variable of type {} already exists",
+                    type_name::<V>()
+                );
+            }
+            Entry::Vacant(e) => {
+                e.insert(Box::new(val));
+            }
+        }
+        self
+    }
 }
 
 /// Trait for accessing and modifying the state of a Vortex session.