feat[array]: validate array encodings on write

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-array/src/lib.rs

@@ -50,6 +50,7 @@ pub mod mask;
 mod mask_future;
 pub mod matchers;
 mod metadata;
+pub mod normalize;
 pub mod optimizer;
 mod partial_ord;
 pub mod patches;

vortex-array/src/normalize.rs

@@ -0,0 +1,52 @@
+// 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 crate::Array;
+use crate::ArrayRef;
+use crate::ArrayVisitor;
+use crate::session::ArrayRegistry;
+
+/// 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,
+    /// The operation to perform when a non-allowed encoding is encountered.
+    pub operation: Operation,
+}
+
+/// The operation to perform when a non-allowed encoding is encountered.
+pub enum Operation {
+    Error,
+    // TODO(joe): add into canonical variant
+}
+
+impl dyn Array + '_ {
+    /// Normalize the array according to given options.
+    ///
+    /// This operation performs a recursive traversal of the array. Any non-allowed encoding is
+    /// normalized per the configured operation.
+    pub fn normalize(self: ArrayRef, 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)
+    }
+
+    fn normalize_with_error(self: &ArrayRef, allowed: &[Id]) -> VortexResult<()> {
+        if !allowed.contains(&self.encoding_id()) {
+            vortex_bail!(AssertionFailed: "normalize forbids encoding ({})", self.encoding_id())
+        }
+
+        for child in ArrayVisitor::children(self) {
+            let child: ArrayRef = child;
+            child.normalize_with_error(allowed)?
+        }
+        Ok(())
+    }
+}

vortex-file/src/strategy.rs

@@ -77,7 +77,7 @@ impl WriteStrategyBuilder {
     /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides
     /// applied.
     pub fn build(self) -> Arc<dyn LayoutStrategy> {
-        let flat = FlatLayoutStrategy::new(self.session);
+        let mut flat = FlatLayoutStrategy::new(self.session);
 
         // 7. for each chunk create a flat layout
         let chunked = ChunkedLayoutStrategy::new(flat.clone());

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

@@ -8,7 +8,10 @@ use vortex_array::ArrayContext;
 use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
 use vortex_array::expr::stats::StatsProvider;
+use vortex_array::normalize::NormalizeOptions;
+use vortex_array::normalize::Operation;
 use vortex_array::serde::SerializeOptions;
+use vortex_array::session::ArraySessionExt;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -140,6 +143,11 @@ impl LayoutStrategy for FlatLayoutStrategy {
             _ => {}
         }
 
+        let chunk = chunk.normalize(&mut NormalizeOptions {
+            allowed: self.session.arrays().registry(),
+            operation: Operation::Error,
+        })?;
+
         let buffers = chunk.serialize(
             &ctx,
             &SerializeOptions {
@@ -199,8 +207,10 @@ mod tests {
     use vortex_dtype::FieldNames;
     use vortex_dtype::Nullability;
     use vortex_error::VortexExpect;
+    use vortex_error::VortexResult;
     use vortex_io::runtime::single::block_on;
     use vortex_mask::AllOr;
+    use vortex_mask::Mask;
 
     use crate::LayoutStrategy;
     use crate::layouts::flat::writer::FlatLayoutStrategy;
@@ -386,4 +396,40 @@ mod tests {
             );
         })
     }
+
+    #[test]
+    fn flat_invalid_array_fails() -> VortexResult<()> {
+        block_on(|handle| async {
+            let prim: PrimitiveArray = (0..10).collect();
+            let filter = prim.filter(Mask::from_indices(10, vec![2, 3]))?;
+
+            let ctx = ArrayContext::empty();
+
+            // Write the array into a byte buffer.
+            let (layout, _segments) = {
+                let segments = Arc::new(TestSegments::default());
+                let (ptr, eof) = SequenceId::root().split();
+                let layout = FlatLayoutStrategy::new(SESSION.clone())
+                    .write_stream(
+                        ctx,
+                        segments.clone(),
+                        filter.to_array_stream().sequenced(ptr),
+                        eof,
+                        handle,
+                    )
+                    .await;
+
+                (layout, segments)
+            };
+
+            let err = layout.expect_err("expected error");
+            assert!(
+                err.to_string()
+                    .contains("normalize forbids encoding (vortex.filter)"),
+                "unexpected error: {err}"
+            );
+
+            Ok(())
+        })
+    }
 }