feat: Support retrieving WriterStrategyBuilder from VortexSession (#6)

* Fix session get-or-default (vortex-data#5662)

The comments described this get-or-default, but instead it was a panic

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

* feat: Support retrieving writer strategy builder from vortex session

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>
Co-authored-by: Nicholas Gates <gatesn@users.noreply.github.com>

Author: 2 people

vortex-file/src/strategy.rs

@@ -111,13 +111,25 @@ pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
 /// 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(Clone)]
 pub struct WriteStrategyBuilder {
     compressor: Option<Arc<dyn CompressorPlugin>>,
     row_block_size: usize,
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
     allow_encodings: Option<ArrayRegistry>,
 }
 
+impl std::fmt::Debug for WriteStrategyBuilder {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("WriteStrategyBuilder")
+            .field("compressor", &self.compressor)
+            .field("row_block_size", &self.row_block_size)
+            .field("field_writers", &format!("({} entries)", self.field_writers.len()))
+            .field("allow_encodings", &self.allow_encodings.as_ref().map(|_| "..."))
+            .finish()
+    }
+}
+
 impl Default for WriteStrategyBuilder {
     /// Create a new empty builder. It can be further configured,
     /// and then finally built yielding the [`LayoutStrategy`].

vortex-file/src/writer.rs

@@ -69,10 +69,14 @@ pub struct VortexWriteOptions {
 pub trait WriteOptionsSessionExt: SessionExt {
     /// Create [`VortexWriteOptions`] for writing to a Vortex file.
     fn write_options(&self) -> VortexWriteOptions {
-        let session = self.session();
+        let maybe_write_strategy_builder = self.get_opt::<WriteStrategyBuilder>();
+        let strategy = maybe_write_strategy_builder
+            .map(|opt| opt.clone().build())
+            .unwrap_or_else(|| WriteStrategyBuilder::default().build());
+
         VortexWriteOptions {
-            strategy: WriteStrategyBuilder::default().build(),
-            session,
+            session: self.session(),
+            strategy,
             exclude_dtype: false,
             file_statistics: PRUNING_STATS.to_vec(),
             max_variable_length_statistics_size: 64,
@@ -462,3 +466,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

@@ -29,7 +29,7 @@ use crate::sequence::SequentialStreamExt;
 /// 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>;
 }
 
@@ -41,7 +41,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

@@ -3,20 +3,14 @@
 
 pub mod registry;
 
-use std::any::Any;
-use std::any::TypeId;
-use std::any::type_name;
+use std::any::{Any, TypeId, type_name};
 use std::fmt::Debug;
-use std::hash::BuildHasherDefault;
-use std::hash::Hasher;
-use std::ops::Deref;
-use std::ops::DerefMut;
+use std::hash::{BuildHasherDefault, Hasher};
+use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 
-use dashmap::DashMap;
-use dashmap::Entry;
-use vortex_error::VortexExpect;
-use vortex_error::vortex_panic;
+use dashmap::{DashMap, Entry};
+use vortex_error::{VortexExpect, vortex_panic};
 
 /// A Vortex session encapsulates the set of extensible arrays, layouts, compute functions, dtypes,
 /// etc. that are available for use in a given context.
@@ -52,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.
@@ -88,18 +102,6 @@ impl SessionExt for VortexSession {
 
     /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
     fn get<V: SessionVar + Default>(&self) -> Ref<'_, V> {
-        // NOTE(ngates): we don't use `entry().or_insert_with_key()` here because the DashMap
-        //  would immediately acquire an exclusive write lock.
-        if let Some(v) = self.0.get(&TypeId::of::<V>()) {
-            return Ref(v.map(|v| {
-                (**v)
-                    .as_any()
-                    .downcast_ref::<V>()
-                    .vortex_expect("Type mismatch - this is a bug")
-            }));
-        }
-
-        // If we get here, the value was not present, so we insert the default with a write lock.
         Ref(self
             .0
             .entry(TypeId::of::<V>())