ALP decoding plugin

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

Author: a10y

encodings/alp/src/alp/mod.rs

@@ -18,8 +18,11 @@ mod compress;
 pub(crate) mod compute;
 mod decompress;
 mod ops;
+mod plugin;
 mod rules;
 
+pub(crate) use plugin::ALPPatchedPlugin;
+
 #[cfg(test)]
 mod tests {
     use prost::Message;

encodings/alp/src/alp/plugin.rs

@@ -0,0 +1,226 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! A custom [`ArrayPlugin`] that lets you load in and deserialize an `ALP` array with interior
+//! patches as a `PatchedArray` that wraps a patchless `ALP` 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::ALP;
+use crate::ALPArrayExt;
+use crate::ALPArrayOwnedExt;
+
+/// Custom deserialization plugin that converts an ALP array with interior
+/// patches into a PatchedArray holding an ALP array.
+#[derive(Debug, Clone)]
+pub(crate) struct ALPPatchedPlugin;
+
+impl ArrayPlugin for ALPPatchedPlugin {
+    fn id(&self) -> ArrayId {
+        // We reuse the existing `ALP` ID so that we can take over its
+        // deserialization pathway.
+        ALP::ID
+    }
+
+    fn serialize(
+        &self,
+        array: &ArrayRef,
+        session: &VortexSession,
+    ) -> VortexResult<Option<Vec<u8>>> {
+        // Delegate to ALP's metadata serde
+        ALP.serialize(array, session)
+    }
+
+    fn deserialize(
+        &self,
+        dtype: &DType,
+        len: usize,
+        metadata: &[u8],
+        buffers: &[BufferHandle],
+        children: &dyn ArrayChildren,
+        session: &VortexSession,
+    ) -> VortexResult<ArrayRef> {
+        let alp_array = ALP
+            .deserialize(dtype, len, metadata, buffers, children, session)?
+            .try_downcast::<ALP>()
+            .map_err(|_| vortex_err!("ALP plugin should only deserialize vortex.alp"))?;
+
+        // Check if there are interior patches to externalize.
+        let Some(patches) = alp_array.patches() else {
+            return Ok(alp_array.into_array());
+        };
+
+        // Extract components and create a new ALP array without patches.
+        let (encoded, exponents, _) = alp_array.into_parts();
+
+        let alp_without_patches = ALP::try_new(encoded, exponents, None)?.into_array();
+
+        let patched = Patched::from_array_and_patches(
+            alp_without_patches,
+            &patches,
+            &mut session.create_execution_ctx(),
+        )?;
+
+        Ok(patched.into_array())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::f64::consts::PI;
+    use std::sync::LazyLock;
+
+    use vortex_array::ArrayPlugin;
+    use vortex_array::IntoArray;
+    use vortex_array::arrays::PatchedArray;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::arrays::patched::PatchedArrayExt;
+    use vortex_array::buffer::BufferHandle;
+    use vortex_array::session::ArraySession;
+    use vortex_array::session::ArraySessionExt;
+    use vortex_error::VortexResult;
+    use vortex_error::vortex_err;
+    use vortex_session::VortexSession;
+
+    use super::ALPPatchedPlugin;
+    use crate::ALP;
+    use crate::ALPArray;
+    use crate::ALPArrayExt;
+    use crate::alp_encode;
+
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::empty().with::<ArraySession>();
+        session.arrays().register(ALPPatchedPlugin);
+        session
+    });
+
+    #[test]
+    fn test_decode_alp_patches() -> VortexResult<()> {
+        // Create values where some don't encode cleanly with ALP, causing patches.
+        // PI doesn't encode cleanly.
+        let values: Vec<f64> = (0..100)
+            .map(|i| if i % 4 == 3 { PI } else { i as f64 })
+            .collect();
+
+        let parray = PrimitiveArray::from_iter(values);
+        let alp_encoded = alp_encode(&parray, None)?;
+
+        assert!(
+            alp_encoded.patches().is_some(),
+            "Expected ALP array to have patches"
+        );
+
+        let array = alp_encoded.as_array();
+
+        let metadata = array.metadata(&SESSION)?.unwrap_or_default();
+        let children = array.children();
+        let buffers = array
+            .buffers()
+            .into_iter()
+            .map(BufferHandle::new_host)
+            .collect::<Vec<_>>();
+
+        let deserialized = ALPPatchedPlugin.deserialize(
+            array.dtype(),
+            array.len(),
+            &metadata,
+            &buffers,
+            &children,
+            &SESSION,
+        )?;
+
+        let patched: PatchedArray = deserialized
+            .try_downcast()
+            .map_err(|a| vortex_err!("Expected Patched, got {}", a.encoding_id()))?;
+
+        let inner_alp: ALPArray = patched
+            .base_array()
+            .clone()
+            .try_downcast()
+            .map_err(|a| vortex_err!("Expected inner ALP, got {}", a.encoding_id()))?;
+
+        assert!(
+            inner_alp.patches().is_none(),
+            "Inner ALP should NOT have patches"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn alp_without_patches_stays_alp() -> VortexResult<()> {
+        // Values that encode cleanly without patches.
+        let values: Vec<f64> = (0..100).map(|i| i as f64).collect();
+        let parray = PrimitiveArray::from_iter(values);
+        let alp_encoded = alp_encode(&parray, None)?;
+
+        assert!(
+            alp_encoded.patches().is_none(),
+            "Expected ALP array without patches"
+        );
+
+        let array = alp_encoded.as_array();
+
+        let metadata = array.metadata(&SESSION)?.unwrap_or_default();
+        let children = array.children();
+        let buffers = array
+            .buffers()
+            .into_iter()
+            .map(BufferHandle::new_host)
+            .collect::<Vec<_>>();
+
+        let deserialized = ALPPatchedPlugin.deserialize(
+            array.dtype(),
+            array.len(),
+            &metadata,
+            &buffers,
+            &children,
+            &SESSION,
+        )?;
+
+        let result = deserialized
+            .try_downcast::<ALP>()
+            .map_err(|a| vortex_err!("Expected deserialized ALP, got {}", a.encoding_id()))?;
+
+        assert!(result.patches().is_none(), "Result should not have patches");
+
+        Ok(())
+    }
+
+    #[test]
+    #[should_panic(expected = "index out of bounds")]
+    fn primitive_array_returns_error() {
+        let array = PrimitiveArray::from_iter([1.0f64, 2.0, 3.0]).into_array();
+
+        let metadata = array.metadata(&SESSION).unwrap().unwrap_or_default();
+        let children = array.children();
+        let buffers = array
+            .buffers()
+            .into_iter()
+            .map(BufferHandle::new_host)
+            .collect::<Vec<_>>();
+
+        // This panics because PrimitiveArray has no children and ALP requires encoded child.
+        let _result = ALPPatchedPlugin.deserialize(
+            array.dtype(),
+            array.len(),
+            &metadata,
+            &buffers,
+            &children,
+            &SESSION,
+        );
+    }
+}

encodings/alp/src/lib.rs

@@ -22,14 +22,21 @@ use vortex_array::aggregate_fn::AggregateFnVTable;
 use vortex_array::aggregate_fn::fns::nan_count::NanCount;
 use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
 use vortex_array::session::ArraySessionExt;
+use vortex_fastlanes::USE_EXPERIMENTAL_PATCHES;
 use vortex_session::VortexSession;
 
 mod alp;
 mod alp_rd;
 
 /// Initialize ALP encoding in the given session.
 pub fn initialize(session: &VortexSession) {
-    session.arrays().register(ALP);
+    // If we're using the experimental Patched encoding, register a shim
+    // for ALP with interior patches to decode as Patched array.
+    if *USE_EXPERIMENTAL_PATCHES {
+        session.arrays().register(ALPPatchedPlugin);
+    } else {
+        session.arrays().register(ALP);
+    }
     session.arrays().register(ALPRD);
 
     // Register the ALP-specific NaN count aggregate kernel.