Normalize Execution

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

Author: gatesn

vortex-array/src/array/erased.rs

@@ -383,6 +383,13 @@ impl ArrayRef {
         self.is::<AnyCanonical>()
     }
 
+    /// Returns a new array with all slots replaced.
+    pub fn with_slots(self, slots: Vec<Option<ArrayRef>>) -> VortexResult<ArrayRef> {
+        vortex_ensure!(self.slots().len() == slots.len(), "slot count mismatch");
+        let vtable = self.vtable().clone_boxed();
+        vtable.with_slots(self, slots)
+    }
+
     /// Returns a new array with the slot at `slot_idx` replaced by `replacement`.
     ///
     /// Takes ownership to allow in-place mutation when the refcount is 1.

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,118 @@ 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 super::NormalizeOptions;
+    use super::Operation;
+    use crate::ArrayRef;
+    use crate::ExecutionCtx;
+    use crate::IntoArray;
+    use crate::arrays::PrimitiveArray;
+    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(())
+    }
 }