cleanup

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-btrblocks/src/builder.rs

@@ -81,7 +81,7 @@ pub fn default_excluded() -> HashSet<SchemeId> {
 ///
 /// ```rust
 /// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt};
-/// use vortex_btrblocks::compressor::integer::DictScheme;
+/// use vortex_btrblocks::schemes::integer::DictScheme;
 ///
 /// // Default compressor - all non-excluded schemes allowed.
 /// let compressor = BtrBlocksCompressorBuilder::default().build();

vortex-btrblocks/src/canonical_compressor.rs

@@ -20,7 +20,7 @@ use crate::CascadingCompressor;
 ///
 /// ```rust
 /// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt};
-/// use vortex_btrblocks::compressor::integer::DictScheme;
+/// use vortex_btrblocks::schemes::integer::DictScheme;
 ///
 /// // Default compressor - all schemes allowed.
 /// let compressor = BtrBlocksCompressor::default();

vortex-btrblocks/src/lib.rs

@@ -41,7 +41,7 @@
 //!
 //! ```rust
 //! use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt};
-//! use vortex_btrblocks::compressor::integer::DictScheme;
+//! use vortex_btrblocks::schemes::integer::DictScheme;
 //! use vortex_array::DynArray;
 //!
 //! // Default compressor with all schemes enabled.

vortex-btrblocks/src/schemes/decimal.rs

@@ -70,9 +70,7 @@ impl Scheme for DecimalScheme {
             _ => return Ok(decimal.into_array()),
         };
 
-        let ctx = CompressorContext::default()
-            .descend()
-            .with_scheme(self.id(), 0);
+        let ctx = CompressorContext::default().descend_with_scheme(self.id(), 0);
         let compressed = compressor.compress_canonical(Canonical::Primitive(prim), ctx)?;
 
         DecimalBytePartsArray::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array())

vortex-btrblocks/src/schemes/float.rs

@@ -122,7 +122,7 @@ impl Scheme for ALPScheme {
             return Ok(0.0);
         }
 
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             // ALP does not compress on its own, we need to be able to cascade it with
             // an integer compressor.
             return Ok(0.0);
@@ -148,7 +148,7 @@ impl Scheme for ALPScheme {
         // to keep them linear for easy indexing.
         let compressed_alp_ints = compressor.compress_canonical(
             Canonical::Primitive(alp_ints),
-            ctx.descend().with_scheme(self.id(), 0),
+            ctx.descend_with_scheme(self.id(), 0),
         )?;
 
         let patches = alp.patches().map(compress_patches).transpose()?;
@@ -220,9 +220,9 @@ impl Scheme for NullDominated {
         1
     }
 
-    fn descendant_exclusions(&self) -> &[DescendantExclusion] {
-        &[DescendantExclusion {
-            excluded: IntSparseScheme::ID,
+    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
+        vec![DescendantExclusion {
+            excluded: IntSparseScheme.id(),
             children: ChildSelection::All,
         }]
     }
@@ -234,7 +234,7 @@ impl Scheme for NullDominated {
         ctx: CompressorContext,
     ) -> VortexResult<f64> {
         // Only use `SparseScheme` if we can cascade.
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(0.0);
         }
 
@@ -260,7 +260,7 @@ impl Scheme for NullDominated {
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        assert!(ctx.allowed_cascading > 0);
+        assert!(!ctx.finished_cascading());
 
         let stats = data.float_stats();
 
@@ -271,7 +271,7 @@ impl Scheme for NullDominated {
             let indices = sparse.patches().indices().to_primitive().narrow()?;
             let compressed_indices = compressor.compress_canonical(
                 Canonical::Primitive(indices.to_primitive()),
-                ctx.descend().with_scheme(self.id(), 0),
+                ctx.descend_with_scheme(self.id(), 0),
             )?;
 
             SparseArray::try_new(

vortex-btrblocks/src/schemes/integer.rs

@@ -39,7 +39,6 @@ use crate::CompressorContext;
 use crate::GenerateStatsOptions;
 use crate::Scheme;
 use crate::SchemeExt;
-use crate::SchemeId;
 use crate::compress_patches;
 use crate::estimate_compression_ratio_with_sampling;
 use crate::schemes::rle;
@@ -62,35 +61,14 @@ pub struct BitPackingScheme;
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct SparseScheme;
 
-impl SparseScheme {
-    /// Constant [`SchemeId`] for use in static exclusion rules.
-    pub const ID: SchemeId = SchemeId {
-        name: "vortex.int.sparse",
-    };
-}
-
 /// Run-end encoding with end positions.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct RunEndScheme;
 
-impl RunEndScheme {
-    /// Constant [`SchemeId`] for use in static exclusion rules.
-    pub const ID: SchemeId = SchemeId {
-        name: "vortex.int.runend",
-    };
-}
-
 /// Sequence encoding for sequential patterns.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct SequenceScheme;
 
-impl SequenceScheme {
-    /// Constant [`SchemeId`] for use in static exclusion rules.
-    pub const ID: SchemeId = SchemeId {
-        name: "vortex.int.sequence",
-    };
-}
-
 /// Pco (pcodec) compression for integers.
 #[cfg(feature = "pco")]
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -158,7 +136,7 @@ impl Scheme for FORScheme {
         ctx: CompressorContext,
     ) -> VortexResult<f64> {
         // Only apply if we are not at the leaf.
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(0.0);
         }
 
@@ -219,9 +197,8 @@ impl Scheme for FORScheme {
         // of bitpacking.
         // NOTE: we could delegate in the future if we had another downstream codec that performs
         //  as well.
-        let mut leaf_ctx = ctx.clone();
-        leaf_ctx.allowed_cascading = 0;
-        let mut biased_data = ArrayAndStats::new(biased.into_array(), ctx.stats_options);
+        let leaf_ctx = ctx.clone().as_leaf();
+        let mut biased_data = ArrayAndStats::new(biased.into_array(), ctx.stats_options());
         let compressed = BitPackingScheme.compress(compressor, &mut biased_data, leaf_ctx)?;
 
         let for_compressed = FoRArray::try_new(compressed, for_array.reference_scalar().clone())?;
@@ -246,18 +223,18 @@ impl Scheme for ZigZagScheme {
         1
     }
 
-    fn descendant_exclusions(&self) -> &[DescendantExclusion] {
-        &[
+    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
+        vec![
             DescendantExclusion {
-                excluded: DictScheme::ID,
+                excluded: DictScheme.id(),
                 children: ChildSelection::All,
             },
             DescendantExclusion {
-                excluded: RunEndScheme::ID,
+                excluded: RunEndScheme.id(),
                 children: ChildSelection::All,
             },
             DescendantExclusion {
-                excluded: SparseScheme::ID,
+                excluded: SparseScheme.id(),
                 children: ChildSelection::All,
             },
         ]
@@ -270,7 +247,7 @@ impl Scheme for ZigZagScheme {
         ctx: CompressorContext,
     ) -> VortexResult<f64> {
         // ZigZag is only useful when we cascade it with another encoding.
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(0.0);
         }
 
@@ -302,7 +279,7 @@ impl Scheme for ZigZagScheme {
         let zag = zigzag_encode(stats.source().clone())?;
         let encoded = zag.encoded().to_primitive();
 
-        let child_ctx = ctx.descend().with_scheme(self.id(), 0);
+        let child_ctx = ctx.descend_with_scheme(self.id(), 0);
         let compressed = compressor.compress_canonical(Canonical::Primitive(encoded), child_ctx)?;
 
         tracing::debug!("zigzag output: {}", compressed.encoding_id());
@@ -383,9 +360,9 @@ impl Scheme for SparseScheme {
         2
     }
 
-    fn descendant_exclusions(&self) -> &[DescendantExclusion] {
-        &[DescendantExclusion {
-            excluded: DictScheme::ID,
+    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
+        vec![DescendantExclusion {
+            excluded: DictScheme.id(),
             children: ChildSelection::All,
         }]
     }
@@ -397,7 +374,7 @@ impl Scheme for SparseScheme {
         ctx: CompressorContext,
     ) -> VortexResult<f64> {
         // Only use `SparseScheme` if we can cascade.
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(0.0);
         }
 
@@ -446,7 +423,7 @@ impl Scheme for SparseScheme {
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        assert!(ctx.allowed_cascading > 0);
+        assert!(!ctx.finished_cascading());
 
         let stats = data.integer_stats();
 
@@ -479,15 +456,15 @@ impl Scheme for SparseScheme {
         )?;
 
         if let Some(sparse) = sparse_encoded.as_opt::<Sparse>() {
-            let values_ctx = ctx.clone().descend().with_scheme(self.id(), 0);
+            let values_ctx = ctx.clone().descend_with_scheme(self.id(), 0);
             let compressed_values = compressor.compress_canonical(
                 Canonical::Primitive(sparse.patches().values().to_primitive()),
                 values_ctx,
             )?;
 
             let indices = sparse.patches().indices().to_primitive().narrow()?;
 
-            let indices_ctx = ctx.descend().with_scheme(self.id(), 1);
+            let indices_ctx = ctx.descend_with_scheme(self.id(), 1);
             let compressed_indices =
                 compressor.compress_canonical(Canonical::Primitive(indices), indices_ctx)?;
 
@@ -517,21 +494,21 @@ impl Scheme for RunEndScheme {
         2
     }
 
-    fn descendant_exclusions(&self) -> &[DescendantExclusion] {
-        &[DescendantExclusion {
-            excluded: DictScheme::ID,
+    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
+        vec![DescendantExclusion {
+            excluded: DictScheme.id(),
             children: ChildSelection::All,
         }]
     }
 
-    fn ancestor_exclusions(&self) -> &[AncestorExclusion] {
+    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
         use vortex_compressor::builtins::FloatDictScheme;
 
-        &[
+        vec![
             // Exclude from FloatDict values child (child 0). This replaces the old ALP
             // conditional propagation of float RLE exclusion to integer RunEnd.
             AncestorExclusion {
-                ancestor: FloatDictScheme::ID,
+                ancestor: FloatDictScheme.id(),
                 children: ChildSelection::One(0),
             },
         ]
@@ -550,7 +527,7 @@ impl Scheme for RunEndScheme {
             return Ok(0.0);
         }
 
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(0.0);
         }
 
@@ -564,18 +541,18 @@ impl Scheme for RunEndScheme {
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        assert!(ctx.allowed_cascading > 0);
+        assert!(!ctx.finished_cascading());
 
         let stats = data.integer_stats();
 
         // Run-end encode the ends.
         let (ends, values) = runend_encode(stats.source());
 
-        let values_ctx = ctx.clone().descend().with_scheme(self.id(), 0);
+        let values_ctx = ctx.clone().descend_with_scheme(self.id(), 0);
         let compressed_values = compressor
             .compress_canonical(Canonical::Primitive(values.to_primitive()), values_ctx)?;
 
-        let ends_ctx = ctx.descend().with_scheme(self.id(), 1);
+        let ends_ctx = ctx.descend_with_scheme(self.id(), 1);
         let compressed_ends =
             compressor.compress_canonical(Canonical::Primitive(ends.to_primitive()), ends_ctx)?;
 
@@ -601,24 +578,24 @@ impl Scheme for SequenceScheme {
         is_integer_primitive(canonical)
     }
 
-    fn ancestor_exclusions(&self) -> &[AncestorExclusion] {
+    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
         use vortex_compressor::builtins::FloatDictScheme;
         use vortex_compressor::builtins::StringDictScheme;
 
-        &[
+        vec![
             // Exclude from IntDict codes.
             AncestorExclusion {
-                ancestor: DictScheme::ID,
+                ancestor: DictScheme.id(),
                 children: ChildSelection::All,
             },
             // Exclude from FloatDict codes (child 1).
             AncestorExclusion {
-                ancestor: FloatDictScheme::ID,
+                ancestor: FloatDictScheme.id(),
                 children: ChildSelection::One(1),
             },
             // Exclude from StringDict codes (child 1).
             AncestorExclusion {
-                ancestor: StringDictScheme::ID,
+                ancestor: StringDictScheme.id(),
                 children: ChildSelection::One(1),
             },
         ]

vortex-btrblocks/src/schemes/rle.rs

@@ -92,9 +92,9 @@ impl<C: RLEConfig> Scheme for RLEScheme<C> {
         3
     }
 
-    fn descendant_exclusions(&self) -> &[DescendantExclusion] {
-        &[DescendantExclusion {
-            excluded: IntDictScheme::ID,
+    fn descendant_exclusions(&self) -> Vec<DescendantExclusion> {
+        vec![DescendantExclusion {
+            excluded: IntDictScheme.id(),
             children: ChildSelection::Many(&[1, 2]),
         }]
     }
@@ -106,7 +106,7 @@ impl<C: RLEConfig> Scheme for RLEScheme<C> {
         ctx: CompressorContext,
     ) -> VortexResult<f64> {
         // RLE is only useful when we cascade it with another encoding.
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(0.0);
         }
 
@@ -137,24 +137,24 @@ impl<C: RLEConfig> Scheme for RLEScheme<C> {
         let stats = data.get_or_insert_with::<C::Stats>(|| C::generate_stats(&array));
         let rle_array = RLEArray::encode(RLEStats::source(stats))?;
 
-        if ctx.allowed_cascading == 0 {
+        if ctx.finished_cascading() {
             return Ok(rle_array.into_array());
         }
 
         let compressed_values = C::compress_values(
             compressor,
             &rle_array.values().to_primitive(),
-            ctx.clone().descend().with_scheme(self.id(), 0),
+            ctx.clone().descend_with_scheme(self.id(), 0),
         )?;
 
         let compressed_indices = compressor.compress_canonical(
             Canonical::Primitive(rle_array.indices().to_primitive().narrow()?),
-            ctx.clone().descend().with_scheme(self.id(), 1),
+            ctx.clone().descend_with_scheme(self.id(), 1),
         )?;
 
         let compressed_offsets = compressor.compress_canonical(
             Canonical::Primitive(rle_array.values_idx_offsets().to_primitive().narrow()?),
-            ctx.descend().with_scheme(self.id(), 2),
+            ctx.descend_with_scheme(self.id(), 2),
         )?;
 
         // SAFETY: Recursive compression doesn't affect the invariants.