Normalize Execution

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

Author: gatesn

vortex-file/src/strategy.rs

@@ -10,6 +10,8 @@ use vortex_alp::ALP;
 // Compressed encodings from encoding crates
 // Canonical array encodings from vortex-array
 use vortex_alp::ALPRD;
+use vortex_array::ArrayId;
+use vortex_array::VTable;
 use vortex_array::arrays::Bool;
 use vortex_array::arrays::Chunked;
 use vortex_array::arrays::Constant;
@@ -26,8 +28,6 @@ use vortex_array::arrays::Struct;
 use vortex_array::arrays::VarBin;
 use vortex_array::arrays::VarBinView;
 use vortex_array::dtype::FieldPath;
-use vortex_array::session::ArrayRegistry;
-use vortex_array::session::ArraySession;
 use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_btrblocks::BtrBlocksCompressorBuilder;
 use vortex_btrblocks::SchemeExt;
@@ -67,6 +67,7 @@ use vortex_btrblocks::{
     schemes::integer,
     schemes::string,
 };
+use vortex_utils::aliases::hash_set::HashSet;
 #[cfg(feature = "zstd")]
 use vortex_zstd::Zstd;
 #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
@@ -78,49 +79,49 @@ const ONE_MEG: u64 = 1 << 20;
 ///
 /// This includes all canonical encodings from vortex-array plus all compressed
 /// encodings from the various encoding crates.
-pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
-    let session = ArraySession::empty();
+pub static ALLOWED_ENCODINGS: LazyLock<HashSet<ArrayId>> = LazyLock::new(|| {
+    let mut allowed = HashSet::new();
 
     // Canonical encodings from vortex-array
-    session.register(Null);
-    session.register(Bool);
-    session.register(Primitive);
-    session.register(Decimal);
-    session.register(VarBin);
-    session.register(VarBinView);
-    session.register(List);
-    session.register(ListView);
-    session.register(FixedSizeList);
-    session.register(Struct);
-    session.register(Extension);
-    session.register(Chunked);
-    session.register(Constant);
-    session.register(Masked);
-    session.register(Dict);
+    allowed.insert(Null.id());
+    allowed.insert(Bool.id());
+    allowed.insert(Primitive.id());
+    allowed.insert(Decimal.id());
+    allowed.insert(VarBin.id());
+    allowed.insert(VarBinView.id());
+    allowed.insert(List.id());
+    allowed.insert(ListView.id());
+    allowed.insert(FixedSizeList.id());
+    allowed.insert(Struct.id());
+    allowed.insert(Extension.id());
+    allowed.insert(Chunked.id());
+    allowed.insert(Constant.id());
+    allowed.insert(Masked.id());
+    allowed.insert(Dict.id());
 
     // Compressed encodings from encoding crates
-    session.register(ALP);
-    session.register(ALPRD);
-    session.register(BitPacked);
-    session.register(ByteBool);
-    session.register(DateTimeParts);
-    session.register(DecimalByteParts);
-    session.register(Delta);
-    session.register(FoR);
-    session.register(FSST);
-    session.register(Pco);
-    session.register(RLE);
-    session.register(RunEnd);
-    session.register(Sequence);
-    session.register(Sparse);
-    session.register(ZigZag);
+    allowed.insert(ALP.id());
+    allowed.insert(ALPRD.id());
+    allowed.insert(BitPacked.id());
+    allowed.insert(ByteBool.id());
+    allowed.insert(DateTimeParts.id());
+    allowed.insert(DecimalByteParts.id());
+    allowed.insert(Delta.id());
+    allowed.insert(FoR.id());
+    allowed.insert(FSST.id());
+    allowed.insert(Pco.id());
+    allowed.insert(RLE.id());
+    allowed.insert(RunEnd.id());
+    allowed.insert(Sequence.id());
+    allowed.insert(Sparse.id());
+    allowed.insert(ZigZag.id());
 
     #[cfg(feature = "zstd")]
-    session.register(Zstd);
+    allowed.insert(Zstd.id());
     #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
-    session.register(ZstdBuffers);
+    allowed.insert(ZstdBuffers.id());
 
-    session.registry().clone()
+    allowed
 });
 
 /// Build a new [writer strategy][LayoutStrategy] to compress and reorganize chunks of a Vortex file.
@@ -132,7 +133,7 @@ pub struct WriteStrategyBuilder {
     compressor_override: Option<Arc<dyn CompressorPlugin>>,
     row_block_size: usize,
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
-    allow_encodings: Option<ArrayRegistry>,
+    allow_encodings: Option<HashSet<ArrayId>>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
 }
 
@@ -169,7 +170,7 @@ impl WriteStrategyBuilder {
     }
 
     /// Override the allowed array encodings for normalization.
-    pub fn with_allow_encodings(mut self, allow_encodings: ArrayRegistry) -> Self {
+    pub fn with_allow_encodings(mut self, allow_encodings: HashSet<ArrayId>) -> Self {
         self.allow_encodings = Some(allow_encodings);
         self
     }

vortex-layout/src/layouts/flat/writer.rs

@@ -4,6 +4,7 @@
 use async_trait::async_trait;
 use futures::StreamExt;
 use vortex_array::ArrayContext;
+use vortex_array::ArrayId;
 use vortex_array::dtype::DType;
 use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
@@ -15,7 +16,6 @@ use vortex_array::scalar::ScalarTruncation;
 use vortex_array::scalar::lower_bound;
 use vortex_array::scalar::upper_bound;
 use vortex_array::serde::SerializeOptions;
-use vortex_array::session::ArrayRegistry;
 use vortex_array::stats::StatsSetRef;
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
@@ -24,6 +24,7 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_io::runtime::Handle;
 use vortex_session::registry::ReadContext;
+use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::IntoLayout;
 use crate::LayoutRef;
@@ -42,7 +43,7 @@ pub struct FlatLayoutStrategy {
     pub max_variable_length_statistics_size: usize,
     /// Optional set of allowed array encodings for normalization.
     /// If None, then all are allowed.
-    pub allowed_encodings: Option<ArrayRegistry>,
+    pub allowed_encodings: Option<HashSet<ArrayId>>,
 }
 
 impl Default for FlatLayoutStrategy {
@@ -69,7 +70,7 @@ impl FlatLayoutStrategy {
     }
 
     /// Set the allowed array encodings for normalization.
-    pub fn with_allow_encodings(mut self, allow_encodings: ArrayRegistry) -> Self {
+    pub fn with_allow_encodings(mut self, allow_encodings: HashSet<ArrayId>) -> Self {
         self.allowed_encodings = Some(allow_encodings);
         self
     }
@@ -226,6 +227,7 @@ mod tests {
     use vortex_io::runtime::single::block_on;
     use vortex_mask::AllOr;
     use vortex_mask::Mask;
+    use vortex_utils::aliases::hash_set::HashSet;
 
     use crate::LayoutStrategy;
     use crate::layouts::flat::writer::FlatLayoutStrategy;
@@ -424,9 +426,8 @@ mod tests {
             let (layout, _segments) = {
                 let segments = Arc::new(TestSegments::default());
                 let (ptr, eof) = SequenceId::root().split();
-                // Only allow primitive encodings - filter arrays should fail.
-                let allowed = ArrayRegistry::default();
-                allowed.register(Primitive::ID, Arc::new(Primitive) as DynVTableRef);
+                // Only allow canonical encodings - filter arrays should fail.
+                let allowed = HashSet::default();
                 let layout = FlatLayoutStrategy::default()
                     .with_allow_encodings(allowed)
                     .write_stream(
@@ -466,9 +467,8 @@ mod tests {
                 let segments = Arc::new(TestSegments::default());
                 let (ptr, eof) = SequenceId::root().split();
                 // Only allow primitive encodings - filter arrays should fail.
-                let allowed = ArrayRegistry::default();
-                allowed.register(Primitive.id(), Arc::new(Primitive) as DynVTableRef);
-                allowed.register(Dict.id(), Arc::new(Dict) as DynVTableRef);
+                let mut allowed = HashSet::default();
+                allowed.insert(Dict.id());
                 let layout = FlatLayoutStrategy::default()
                     .with_allow_encodings(allowed)
                     .write_stream(