feat[layout]: extract a LayoutStrategyEncodingValidator from the FlatLayoutStrategy

vortex-cuda/src/layout.rs

@@ -14,7 +14,6 @@ use futures::FutureExt;
 use futures::StreamExt;
 use futures::future::BoxFuture;
 use vortex::array::ArrayContext;
-use vortex::array::ArrayId;
 use vortex::array::ArrayRef;
 use vortex::array::ArrayVTable;
 use vortex::array::DeserializeMetadata;
@@ -26,8 +25,6 @@ use vortex::array::expr::Expression;
 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::serde::SerializedArray;
 use vortex::array::stats::StatsSetRef;
@@ -68,7 +65,6 @@ use vortex::session::VortexSession;
 use vortex::session::registry::CachedId;
 use vortex::session::registry::ReadContext;
 use vortex::utils::aliases::hash_map::HashMap;
-use vortex::utils::aliases::hash_set::HashSet;
 
 /// A buffer inlined into layout metadata for host-side access.
 #[derive(Clone, prost::Message)]
@@ -403,16 +399,13 @@ pub struct CudaFlatLayoutStrategy {
     pub include_padding: bool,
     /// Maximum length of variable length statistics.
     pub max_variable_length_statistics_size: usize,
-    /// Optional set of allowed array encodings for normalization.
-    pub allowed_encodings: Option<HashSet<ArrayId>>,
 }
 
 impl Default for CudaFlatLayoutStrategy {
     fn default() -> Self {
         Self {
             include_padding: true,
             max_variable_length_statistics_size: 64,
-            allowed_encodings: None,
         }
     }
 }
@@ -427,11 +420,6 @@ impl CudaFlatLayoutStrategy {
         self.max_variable_length_statistics_size = size;
         self
     }
-
-    pub fn with_allow_encodings(mut self, allow_encodings: HashSet<ArrayId>) -> Self {
-        self.allowed_encodings = Some(allow_encodings);
-        self
-    }
 }
 
 fn truncate_scalar_stat<F: Fn(Scalar) -> Option<(Scalar, bool)>>(
@@ -508,15 +496,6 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
             _ => {}
         }
 
-        let chunk = if let Some(allowed) = &options.allowed_encodings {
-            chunk.normalize(&mut NormalizeOptions {
-                allowed,
-                operation: Operation::Error,
-            })?
-        } else {
-            chunk
-        };
-
         // Scan for constant array buffers before serialization (while data is still on host).
         let host_buffers = extract_constant_buffers(&chunk);
 

vortex-file/src/strategy.rs

@@ -43,6 +43,7 @@ use vortex_fastlanes::FoR;
 use vortex_fastlanes::RLE;
 use vortex_fsst::FSST;
 use vortex_layout::LayoutStrategy;
+use vortex_layout::LayoutStrategyEncodingValidator;
 use vortex_layout::layouts::buffered::BufferedStrategy;
 use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy;
 use vortex_layout::layouts::collect::CollectStrategy;
@@ -217,8 +218,8 @@ impl WriteStrategyBuilder {
 
     /// Override the allowed array encodings for normalization.
     ///
-    /// The flat leaf writer uses this set when deciding whether an existing encoded array can be
-    /// written as-is or must be normalized before serialization.
+    /// The configured flat leaf strategy is wrapped in a [`LayoutStrategyEncodingValidator`]
+    /// that recursively checks every chunk before passing it to the leaf writer.
     pub fn with_allow_encodings(mut self, allow_encodings: HashSet<ArrayId>) -> Self {
         self.allow_encodings = Some(allow_encodings);
         self
@@ -262,11 +263,14 @@ impl WriteStrategyBuilder {
     pub fn build(self) -> Arc<dyn LayoutStrategy> {
         let flat: Arc<dyn LayoutStrategy> = if let Some(flat) = self.flat_strategy {
             flat
-        } else if let Some(allow_encodings) = self.allow_encodings {
-            Arc::new(FlatLayoutStrategy::default().with_allow_encodings(allow_encodings))
         } else {
             Arc::new(FlatLayoutStrategy::default())
         };
+        let flat: Arc<dyn LayoutStrategy> = if let Some(allow_encodings) = self.allow_encodings {
+            Arc::new(LayoutStrategyEncodingValidator::new(flat, allow_encodings))
+        } else {
+            flat
+        };
 
         // 7. for each chunk create a flat layout
         let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat));

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

@@ -4,13 +4,10 @@
 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;
 use vortex_array::expr::stats::StatsProvider;
-use vortex_array::normalize::NormalizeOptions;
-use vortex_array::normalize::Operation;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar::ScalarTruncation;
 use vortex_array::scalar::lower_bound;
@@ -24,7 +21,6 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_session::VortexSession;
 use vortex_session::registry::ReadContext;
-use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::IntoLayout;
 use crate::LayoutRef;
@@ -43,17 +39,13 @@ pub struct FlatLayoutStrategy {
     pub include_padding: bool,
     /// Maximum length of variable length statistics
     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<HashSet<ArrayId>>,
 }
 
 impl Default for FlatLayoutStrategy {
     fn default() -> Self {
         Self {
             include_padding: true,
             max_variable_length_statistics_size: 64,
-            allowed_encodings: None,
         }
     }
 }
@@ -70,12 +62,6 @@ impl FlatLayoutStrategy {
         self.max_variable_length_statistics_size = size;
         self
     }
-
-    /// Set the allowed array encodings for normalization.
-    pub fn with_allow_encodings(mut self, allow_encodings: HashSet<ArrayId>) -> Self {
-        self.allowed_encodings = Some(allow_encodings);
-        self
-    }
 }
 
 fn truncate_scalar_stat<F: Fn(Scalar) -> Option<(Scalar, bool)>>(
@@ -158,15 +144,6 @@ impl LayoutStrategy for FlatLayoutStrategy {
             _ => {}
         }
 
-        let chunk = if let Some(allowed) = &self.allowed_encodings {
-            chunk.normalize(&mut NormalizeOptions {
-                allowed,
-                operation: Operation::Error,
-            })?
-        } else {
-            chunk
-        };
-
         let buffers = chunk.serialize(
             &ctx,
             session,
@@ -239,6 +216,7 @@ mod tests {
     use vortex_utils::aliases::hash_set::HashSet;
 
     use crate::LayoutStrategy;
+    use crate::LayoutStrategyEncodingValidator;
     use crate::layouts::flat::writer::FlatLayoutStrategy;
     use crate::segments::TestSegments;
     use crate::sequence::SequenceId;
@@ -449,16 +427,16 @@ mod tests {
                 let (ptr, eof) = SequenceId::root().split();
                 // Disallow all encodings so filter arrays fail normalization immediately.
                 let allowed = HashSet::default();
-                let layout = FlatLayoutStrategy::default()
-                    .with_allow_encodings(allowed)
-                    .write_stream(
-                        ctx,
-                        Arc::<TestSegments>::clone(&segments),
-                        filter.into_array().to_array_stream().sequenced(ptr),
-                        eof,
-                        &session,
-                    )
-                    .await;
+                let layout =
+                    LayoutStrategyEncodingValidator::new(FlatLayoutStrategy::default(), allowed)
+                        .write_stream(
+                            ctx,
+                            Arc::<TestSegments>::clone(&segments),
+                            filter.into_array().to_array_stream().sequenced(ptr),
+                            eof,
+                            &session,
+                        )
+                        .await;
 
                 (layout, segments)
             };
@@ -491,16 +469,16 @@ mod tests {
                 // Only allow the dict encoding; canonical primitive children remain permitted.
                 let mut allowed = HashSet::default();
                 allowed.insert(Dict.id());
-                let layout = FlatLayoutStrategy::default()
-                    .with_allow_encodings(allowed)
-                    .write_stream(
-                        ctx,
-                        Arc::<TestSegments>::clone(&segments),
-                        dict.into_array().to_array_stream().sequenced(ptr),
-                        eof,
-                        &session,
-                    )
-                    .await;
+                let layout =
+                    LayoutStrategyEncodingValidator::new(FlatLayoutStrategy::default(), allowed)
+                        .write_stream(
+                            ctx,
+                            Arc::<TestSegments>::clone(&segments),
+                            dict.into_array().to_array_stream().sequenced(ptr),
+                            eof,
+                            &session,
+                        )
+                        .await;
 
                 (layout, segments)
             };

vortex-layout/src/strategy.rs

@@ -1,15 +1,24 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::sync::Arc;
+
 use async_trait::async_trait;
+use futures::StreamExt;
 use vortex_array::ArrayContext;
+use vortex_array::ArrayId;
+use vortex_array::normalize::NormalizeOptions;
+use vortex_array::normalize::Operation;
 use vortex_error::VortexResult;
 use vortex_session::VortexSession;
+use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::LayoutRef;
 use crate::segments::SegmentSinkRef;
 use crate::sequence::SendableSequentialStream;
 use crate::sequence::SequencePointer;
+use crate::sequence::SequentialStreamAdapter;
+use crate::sequence::SequentialStreamExt;
 
 // [layout writer]
 /// Writes an ordered array stream into a layout tree and segment sink.
@@ -71,8 +80,65 @@ pub trait LayoutStrategy: 'static + Send + Sync {
     }
 }
 
+/// A layout strategy wrapper that rejects arrays containing encodings outside an allow-list.
+///
+/// Canonical encodings are always permitted. Every chunk is recursively validated before it is
+/// passed to the wrapped strategy.
+#[derive(Clone)]
+pub struct LayoutStrategyEncodingValidator {
+    child: Arc<dyn LayoutStrategy>,
+    allowed_encodings: Arc<HashSet<ArrayId>>,
+}
+
+impl LayoutStrategyEncodingValidator {
+    /// Creates a validator around `child` using the supplied encoding allow-list.
+    pub fn new<S: LayoutStrategy>(child: S, allowed_encodings: HashSet<ArrayId>) -> Self {
+        Self {
+            child: Arc::new(child),
+            allowed_encodings: Arc::new(allowed_encodings),
+        }
+    }
+}
+
+#[async_trait]
+impl LayoutStrategy for LayoutStrategyEncodingValidator {
+    async fn write_stream(
+        &self,
+        ctx: ArrayContext,
+        segment_sink: SegmentSinkRef,
+        stream: SendableSequentialStream,
+        eof: SequencePointer,
+        session: &VortexSession,
+    ) -> VortexResult<LayoutRef> {
+        let dtype = stream.dtype().clone();
+        let allowed_encodings = Arc::clone(&self.allowed_encodings);
+        let stream = stream.map(move |chunk| {
+            let (sequence_id, chunk) = chunk?;
+            let chunk = chunk.normalize(&mut NormalizeOptions {
+                allowed: &allowed_encodings,
+                operation: Operation::Error,
+            })?;
+            Ok((sequence_id, chunk))
+        });
+
+        self.child
+            .write_stream(
+                ctx,
+                segment_sink,
+                SequentialStreamAdapter::new(dtype, stream).sendable(),
+                eof,
+                session,
+            )
+            .await
+    }
+
+    fn buffered_bytes(&self) -> u64 {
+        self.child.buffered_bytes()
+    }
+}
+
 #[async_trait]
-impl LayoutStrategy for std::sync::Arc<dyn LayoutStrategy> {
+impl LayoutStrategy for Arc<dyn LayoutStrategy> {
     async fn write_stream(
         &self,
         ctx: ArrayContext,