saev

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

encodings/fastlanes/src/bitpacking/mod.rs

@@ -11,6 +11,9 @@ pub use array::unpack_iter;
 
 pub(crate) mod compute;
 
+mod plugin;
 mod vtable;
+
+pub use plugin::BitPackedPatchedPlugin;
 pub use vtable::BitPacked;
 pub use vtable::BitPackedArray;

encodings/fastlanes/src/bitpacking/plugin.rs

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! A custom [`ArrayPlugin`] that lets you load in and deserialize a `BitPacked` array with interior
+//! patches as a `PatchedArray` that wraps a patchless `BitPacked` array.
+//!
+//! This enables zero-cost backward compatibility with previously written datasets.
+
+use vortex_array::ArrayId;
+use vortex_array::ArrayPlugin;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::Patched;
+use vortex_array::buffer::BufferHandle;
+use vortex_array::dtype::DType;
+use vortex_array::serde::ArrayChildren;
+use vortex_error::VortexResult;
+use vortex_error::vortex_err;
+use vortex_session::VortexSession;
+
+use crate::BitPacked;
+use crate::BitPackedArrayExt;
+
+/// Custom deserialization plugin that converts a BitPacked array with interior
+/// Patches into a PatchedArray holding a BitPacked array.
+#[derive(Debug, Clone)]
+pub struct BitPackedPatchedPlugin;
+
+impl ArrayPlugin for BitPackedPatchedPlugin {
+    fn id(&self) -> ArrayId {
+        // We reuse the existing `BitPacked` ID so that we can take over its
+        // deserialization pathway.
+        BitPacked::ID
+    }
+
+    fn deserialize(
+        &self,
+        dtype: &DType,
+        len: usize,
+        metadata: &[u8],
+        buffers: &[BufferHandle],
+        children: &dyn ArrayChildren,
+        session: &VortexSession,
+    ) -> VortexResult<ArrayRef> {
+        let bitpacked = BitPacked
+            .deserialize(dtype, len, metadata, buffers, children, session)?
+            .try_downcast::<BitPacked>()
+            .map_err(|_| {
+                vortex_err!("BitPacked plugin should only deserialize vortex.bitpacked")
+            })?;
+
+        // Create a new BitPackedArray without the interior patches installed.
+        let Some(patches) = bitpacked.patches() else {
+            return Ok(bitpacked.into_array());
+        };
+
+        let packed = bitpacked.packed().clone();
+        let ptype = bitpacked.dtype().as_ptype();
+        let validity = bitpacked.validity()?;
+        let bw = bitpacked.bit_width;
+        let len = bitpacked.len();
+        let offset = bitpacked.offset();
+
+        let bitpacked_without_patches =
+            BitPacked::try_new(packed, ptype, validity, None, bw, len, offset)?.into_array();
+
+        let patched = Patched::from_array_and_patches(
+            bitpacked_without_patches,
+            &patches,
+            &mut session.create_execution_ctx(),
+        )?;
+
+        Ok(patched.into_array())
+    }
+}

vortex-file/src/strategy.rs

@@ -3,14 +3,14 @@
 
 //! This module defines the default layout strategy for a Vortex file.
 
+use std::env;
 use std::sync::Arc;
 use std::sync::LazyLock;
 
 use vortex_alp::ALP;
 // Compressed encodings from encoding crates
 // Canonical array encodings from vortex-array
 use vortex_alp::ALPRD;
-use vortex_array::arrays::Bool;
 use vortex_array::arrays::Chunked;
 use vortex_array::arrays::Constant;
 use vortex_array::arrays::Decimal;
@@ -25,6 +25,7 @@ use vortex_array::arrays::Primitive;
 use vortex_array::arrays::Struct;
 use vortex_array::arrays::VarBin;
 use vortex_array::arrays::VarBinView;
+use vortex_array::arrays::{Bool, Patched};
 use vortex_array::dtype::FieldPath;
 use vortex_array::session::ArrayRegistry;
 use vortex_array::session::ArraySession;
@@ -35,6 +36,7 @@ use vortex_bytebool::ByteBool;
 use vortex_datetime_parts::DateTimeParts;
 use vortex_decimal_byte_parts::DecimalByteParts;
 use vortex_fastlanes::BitPacked;
+use vortex_fastlanes::BitPackedPatchedPlugin;
 use vortex_fastlanes::Delta;
 use vortex_fastlanes::FoR;
 use vortex_fastlanes::RLE;
@@ -67,6 +69,10 @@ use vortex_zstd::ZstdBuffers;
 
 const ONE_MEG: u64 = 1 << 20;
 
+/// Check if we're using experimental patches deserialization
+static USE_EXPERIMENTAL_PATCHES: LazyLock<bool> =
+    LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_PATCHED_ARRAY").is_ok());
+
 /// Static registry of all allowed array encodings for file writing.
 ///
 /// This includes all canonical encodings from vortex-array plus all compressed
@@ -94,7 +100,14 @@ pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
     // Compressed encodings from encoding crates
     session.register(ALP);
     session.register(ALPRD);
-    session.register(BitPacked);
+    if *USE_EXPERIMENTAL_PATCHES {
+        session.register(Patched);
+        // Deserialize BitPacked arrays with Patches as PatchedArray of BitPacked
+        // without the interior patches.
+        session.register(BitPackedPatchedPlugin)
+    } else {
+        session.register(BitPacked);
+    }
     session.register(ByteBool);
     session.register(DateTimeParts);
     session.register(DecimalByteParts);