Normalize Execution

vortex-array/src/normalize.rs

@@ -1,27 +1,27 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use itertools::Itertools;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_session::registry::Id;
+use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::ArrayRef;
-use crate::session::ArrayRegistry;
+use crate::ExecutionCtx;
 
 /// Options for normalizing an array.
 pub struct NormalizeOptions<'a> {
     /// The set of allowed array encodings (in addition to the canonical ones) that are permitted
     /// in the normalized array.
-    pub allowed: &'a ArrayRegistry,
+    pub allowed: &'a HashSet<Id>,
     /// The operation to perform when a non-allowed encoding is encountered.
-    pub operation: Operation,
+    pub operation: Operation<'a>,
 }
 
 /// The operation to perform when a non-allowed encoding is encountered.
-pub enum Operation {
+pub enum Operation<'a> {
     Error,
-    // TODO(joe): add into canonical variant
+    Execute(&'a mut ExecutionCtx),
 }
 
 impl ArrayRef {
@@ -30,13 +30,17 @@ impl ArrayRef {
     /// This operation performs a recursive traversal of the array. Any non-allowed encoding is
     /// normalized per the configured operation.
     pub fn normalize(self, options: &mut NormalizeOptions) -> VortexResult<ArrayRef> {
-        let array_ids = options.allowed.ids().collect_vec();
-        self.normalize_with_error(&array_ids)?;
-        // Note this takes ownership so we can at a later date remove non-allowed encodings.
-        Ok(self)
+        match &mut options.operation {
+            Operation::Error => {
+                self.normalize_with_error(options.allowed)?;
+                // Note this takes ownership so we can at a later date remove non-allowed encodings.
+                Ok(self)
+            }
+            Operation::Execute(ctx) => self.normalize_with_execution(options.allowed, ctx),
+        }
     }
 
-    fn normalize_with_error(&self, allowed: &[Id]) -> VortexResult<()> {
+    fn normalize_with_error(&self, allowed: &HashSet<Id>) -> VortexResult<()> {
         if !allowed.contains(&self.encoding_id()) {
             vortex_bail!(AssertionFailed: "normalize forbids encoding ({})", self.encoding_id())
         }
@@ -46,4 +50,158 @@ impl ArrayRef {
         }
         Ok(())
     }
+
+    fn normalize_with_execution(
+        self,
+        allowed: &HashSet<Id>,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let mut normalized = self;
+
+        // Top-first execute the array tree while we hit non-allowed encodings.
+        while !allowed.contains(&normalized.encoding_id()) {
+            normalized = normalized.execute(ctx)?;
+        }
+
+        // Now we've normalized the root, we need to ensure the children are normalized also.
+        let slots = normalized.slots();
+        let mut normalized_slots = Vec::with_capacity(slots.len());
+        let mut any_slot_changed = false;
+
+        for slot in slots {
+            match slot {
+                Some(child) => {
+                    let normalized_child = child.clone().normalize(&mut NormalizeOptions {
+                        allowed,
+                        operation: Operation::Execute(ctx),
+                    })?;
+                    any_slot_changed |= !ArrayRef::ptr_eq(child, &normalized_child);
+                    normalized_slots.push(Some(normalized_child));
+                }
+                None => normalized_slots.push(None),
+            }
+        }
+
+        if any_slot_changed {
+            normalized = normalized.with_slots(normalized_slots)?;
+        }
+
+        Ok(normalized)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
+    use vortex_utils::aliases::hash_set::HashSet;
+
+    use super::NormalizeOptions;
+    use super::Operation;
+    use crate::ArrayRef;
+    use crate::ExecutionCtx;
+    use crate::IntoArray;
+    use crate::arrays::Dict;
+    use crate::arrays::DictArray;
+    use crate::arrays::Primitive;
+    use crate::arrays::PrimitiveArray;
+    use crate::arrays::Slice;
+    use crate::arrays::SliceArray;
+    use crate::arrays::StructArray;
+    use crate::assert_arrays_eq;
+    use crate::validity::Validity;
+
+    #[test]
+    fn normalize_with_execution_keeps_parent_when_children_are_unchanged() -> VortexResult<()> {
+        let field = PrimitiveArray::from_iter(0i32..4).into_array();
+        let array = StructArray::try_new(
+            ["field"].into(),
+            vec![field.clone()],
+            field.len(),
+            Validity::NonNullable,
+        )?
+        .into_array();
+        let allowed = HashSet::from_iter([array.encoding_id(), field.encoding_id()]);
+        let mut ctx = ExecutionCtx::new(VortexSession::empty());
+
+        let normalized = array.clone().normalize(&mut NormalizeOptions {
+            allowed: &allowed,
+            operation: Operation::Execute(&mut ctx),
+        })?;
+
+        assert!(ArrayRef::ptr_eq(&array, &normalized));
+        Ok(())
+    }
+
+    #[test]
+    fn normalize_with_execution_rebuilds_parent_when_a_child_changes() -> VortexResult<()> {
+        let unchanged = PrimitiveArray::from_iter(0i32..4).into_array();
+        let sliced =
+            SliceArray::new(PrimitiveArray::from_iter(10i32..20).into_array(), 2..6).into_array();
+        let array = StructArray::try_new(
+            ["lhs", "rhs"].into(),
+            vec![unchanged.clone(), sliced.clone()],
+            unchanged.len(),
+            Validity::NonNullable,
+        )?
+        .into_array();
+        let allowed = HashSet::from_iter([array.encoding_id(), unchanged.encoding_id()]);
+        let mut ctx = ExecutionCtx::new(VortexSession::empty());
+
+        let normalized = array.clone().normalize(&mut NormalizeOptions {
+            allowed: &allowed,
+            operation: Operation::Execute(&mut ctx),
+        })?;
+
+        assert!(!ArrayRef::ptr_eq(&array, &normalized));
+
+        let original_children = array.children();
+        let normalized_children = normalized.children();
+        assert!(ArrayRef::ptr_eq(
+            &original_children[0],
+            &normalized_children[0]
+        ));
+        assert!(!ArrayRef::ptr_eq(
+            &original_children[1],
+            &normalized_children[1]
+        ));
+        assert_arrays_eq!(normalized_children[1], PrimitiveArray::from_iter(12i32..16));
+
+        Ok(())
+    }
+
+    #[test]
+    fn normalize_slice_of_dict_returns_dict() -> VortexResult<()> {
+        let codes = PrimitiveArray::from_iter(vec![0u32, 1, 0, 1, 2]).into_array();
+        let values = PrimitiveArray::from_iter(vec![10i32, 20, 30]).into_array();
+        let dict = DictArray::try_new(codes, values)?.into_array();
+
+        // Slice the dict array to get a SliceArray wrapping a DictArray.
+        let sliced = SliceArray::new(dict, 1..4).into_array();
+        assert_eq!(sliced.encoding_id(), Slice::ID);
+
+        let allowed = HashSet::from_iter([Dict::ID, Primitive::ID]);
+        let mut ctx = ExecutionCtx::new(VortexSession::empty());
+
+        println!("sliced {}", sliced.display_tree());
+
+        let normalized = sliced.normalize(&mut NormalizeOptions {
+            allowed: &allowed,
+            operation: Operation::Execute(&mut ctx),
+        })?;
+
+        println!("after {}", normalized.display_tree());
+
+        // The normalized result should be a DictArray, not a SliceArray.
+        assert_eq!(normalized.encoding_id(), Dict::ID);
+        assert_eq!(normalized.len(), 3);
+
+        // Verify the data: codes [1,0,1] -> values [20, 10, 20]
+        assert_arrays_eq!(
+            normalized.to_canonical()?,
+            PrimitiveArray::from_iter(vec![20i32, 10, 20])
+        );
+
+        Ok(())
+    }
 }

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::BtrBlocksCompressorBuilder;
 use vortex_btrblocks::SchemeExt;
 use vortex_btrblocks::schemes::integer::IntDictScheme;
@@ -60,6 +60,9 @@ use vortex_sparse::Sparse;
 use vortex_tensor::encodings::turboquant::TurboQuant;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_zigzag::ZigZag;
+
+
+use vortex_utils::aliases::hash_set::HashSet;
 #[cfg(feature = "zstd")]
 use vortex_zstd::Zstd;
 #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
@@ -71,51 +74,51 @@ 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);
+    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());
     #[cfg(feature = "unstable_encodings")]
-    session.register(TurboQuant);
-    session.register(ZigZag);
+    allowed.insert(TurboQuant.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
 });
 
 /// How the compressor was configured on [`WriteStrategyBuilder`].
@@ -138,7 +141,7 @@ pub struct WriteStrategyBuilder {
     compressor: CompressorConfig,
     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>>,
 }
 
@@ -175,7 +178,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
     }