From 9c7cddc4fb2a15d844378aa369c30b873e55fb66 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 18:36:14 +0100 Subject: [PATCH] feat[layout]: extract a `LayoutStrategyEncodingValidator` from the flat layout to allow composing encoding validation with different layout strategies This will be used to adding edition encoding filtering Signed-off-by: Joe Isaacs --- vortex-cuda/src/layout.rs | 21 -------- vortex-file/src/strategy.rs | 12 +++-- vortex-layout/src/layouts/flat/writer.rs | 64 ++++++++-------------- vortex-layout/src/strategy.rs | 68 +++++++++++++++++++++++- 4 files changed, 96 insertions(+), 69 deletions(-) diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 1b1d37cd14d..6bd81b59f03 100644 --- a/vortex-cuda/src/layout.rs +++ b/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,8 +399,6 @@ 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>, } impl Default for CudaFlatLayoutStrategy { @@ -412,7 +406,6 @@ impl Default for CudaFlatLayoutStrategy { 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) -> Self { - self.allowed_encodings = Some(allow_encodings); - self - } } fn truncate_scalar_stat 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); diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 3cedbc1d68c..763db758cb7 100644 --- a/vortex-file/src/strategy.rs +++ b/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) -> Self { self.allow_encodings = Some(allow_encodings); self @@ -262,11 +263,14 @@ impl WriteStrategyBuilder { pub fn build(self) -> Arc { let flat: Arc = 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 = 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)); diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index f6b6d5acf02..55c07275301 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/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,9 +39,6 @@ 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>, } impl Default for FlatLayoutStrategy { @@ -53,7 +46,6 @@ impl Default for FlatLayoutStrategy { 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) -> Self { - self.allowed_encodings = Some(allow_encodings); - self - } } fn truncate_scalar_stat 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::::clone(&segments), - filter.into_array().to_array_stream().sequenced(ptr), - eof, - &session, - ) - .await; + let layout = + LayoutStrategyEncodingValidator::new(FlatLayoutStrategy::default(), allowed) + .write_stream( + ctx, + Arc::::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::::clone(&segments), - dict.into_array().to_array_stream().sequenced(ptr), - eof, - &session, - ) - .await; + let layout = + LayoutStrategyEncodingValidator::new(FlatLayoutStrategy::default(), allowed) + .write_stream( + ctx, + Arc::::clone(&segments), + dict.into_array().to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await; (layout, segments) }; diff --git a/vortex-layout/src/strategy.rs b/vortex-layout/src/strategy.rs index 4fcd24dd02f..f9f91d9d82f 100644 --- a/vortex-layout/src/strategy.rs +++ b/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, + allowed_encodings: Arc>, +} + +impl LayoutStrategyEncodingValidator { + /// Creates a validator around `child` using the supplied encoding allow-list. + pub fn new(child: S, allowed_encodings: HashSet) -> 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 { + 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 { +impl LayoutStrategy for Arc { async fn write_stream( &self, ctx: ArrayContext,