add a LazyPatchedArray

this lets us deserialize BPArray with Patches without eagerly
transposing

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

Author: a10y

vortex-array/src/arrays/lazy_patched/mod.rs

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod vtable;
+
+pub use vtable::*;

vortex-array/src/arrays/lazy_patched/vtable/mod.rs

@@ -0,0 +1,195 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod operations;
+mod validity;
+
+use std::hash::Hasher;
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_ensure_eq;
+use vortex_error::vortex_panic;
+use vortex_session::VortexSession;
+
+use crate::ArrayEq;
+use crate::ArrayHash;
+use crate::ArrayRef;
+use crate::DeserializeMetadata;
+use crate::DynArray;
+use crate::ExecutionCtx;
+use crate::ExecutionResult;
+use crate::IntoArray;
+use crate::Precision;
+use crate::ProstMetadata;
+use crate::SerializeMetadata;
+use crate::arrays::PatchedArray;
+use crate::buffer::BufferHandle;
+use crate::dtype::DType;
+use crate::patches::Patches;
+use crate::serde::ArrayChildren;
+use crate::stats::StatsSetRef;
+use crate::vtable;
+use crate::vtable::ArrayId;
+use crate::vtable::VTable;
+use crate::vtable::ValidityVTableFromChild;
+
+#[derive(Clone, Debug)]
+pub struct LazyPatched;
+
+vtable!(LazyPatched);
+
+#[derive(Clone, prost::Message)]
+pub struct LazyPatchedMetadata {
+    #[prost(uint32, tag = "1")]
+    pub(crate) num_patches: u32,
+}
+
+impl VTable for LazyPatched {
+    type Array = LazyPatchedArray;
+    type Metadata = ProstMetadata<LazyPatchedMetadata>;
+
+    type OperationsVTable = Self;
+    type ValidityVTable = ValidityVTableFromChild;
+
+    fn vtable(_array: &Self::Array) -> &Self {
+        &LazyPatched
+    }
+
+    fn id(&self) -> ArrayId {
+        ArrayId::new_ref("vortex.patched_lazy")
+    }
+
+    fn len(array: &Self::Array) -> usize {
+        array.inner.len()
+    }
+
+    fn dtype(array: &Self::Array) -> &DType {
+        array.inner.dtype()
+    }
+
+    fn stats(_array: &Self::Array) -> StatsSetRef<'_> {
+        todo!()
+    }
+
+    fn array_hash<H: Hasher>(array: &Self::Array, state: &mut H, precision: Precision) {
+        array.inner.array_hash(state, precision);
+        array.patches.array_hash(state, precision);
+    }
+
+    fn array_eq(array: &Self::Array, other: &Self::Array, precision: Precision) -> bool {
+        array.inner.array_eq(&other.inner, precision)
+            && array.patches.array_eq(&other.patches, precision)
+    }
+
+    fn nbuffers(_array: &Self::Array) -> usize {
+        0
+    }
+
+    fn buffer(_array: &Self::Array, _idx: usize) -> BufferHandle {
+        vortex_panic!("LazyPatched array holds no buffers")
+    }
+
+    fn buffer_name(_array: &Self::Array, _idx: usize) -> Option<String> {
+        vortex_panic!("LazyPatched array holds no buffers")
+    }
+
+    fn nchildren(_array: &Self::Array) -> usize {
+        3
+    }
+
+    fn child(array: &Self::Array, idx: usize) -> ArrayRef {
+        match idx {
+            0 => array.inner.clone(),
+            1 => array.patches.indices().clone(),
+            2 => array.patches.values().clone(),
+            _ => unreachable!("invalid LazyPatched child index {}", idx),
+        }
+    }
+
+    fn child_name(_array: &Self::Array, idx: usize) -> String {
+        match idx {
+            0 => "inner".to_string(),
+            1 => "patch_indices".to_string(),
+            2 => "patch_values".to_string(),
+            _ => unreachable!("invalid LazyPatched child index {}", idx),
+        }
+    }
+
+    fn metadata(array: &Self::Array) -> VortexResult<Self::Metadata> {
+        let num_patches = u32::try_from(array.patches.num_patches())?;
+
+        Ok(ProstMetadata(LazyPatchedMetadata { num_patches }))
+    }
+
+    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
+        Ok(Some(metadata.serialize()))
+    }
+
+    fn deserialize(
+        bytes: &[u8],
+        _dtype: &DType,
+        _len: usize,
+        _buffers: &[BufferHandle],
+        _session: &VortexSession,
+    ) -> VortexResult<Self::Metadata> {
+        let deserialized = <Self::Metadata>::deserialize(bytes)?;
+        Ok(ProstMetadata(deserialized))
+    }
+
+    fn build(
+        dtype: &DType,
+        len: usize,
+        metadata: &Self::Metadata,
+        _buffers: &[BufferHandle],
+        children: &dyn ArrayChildren,
+    ) -> VortexResult<Self::Array> {
+        // There should be 3 children
+        // 1. inner
+        // 2. patch_indices
+        // 3. patch_values
+        vortex_ensure!(
+            children.len() == 3,
+            "expected exactly 3 children from LazyPatched, found {}",
+            children.len()
+        );
+
+        let inner = children.get(0, dtype, len)?;
+
+        let num_patches = metadata.num_patches as usize;
+        let patch_indices = children.get(1, dtype, num_patches)?;
+        let patch_values = children.get(2, dtype, num_patches)?;
+
+        let patches = Patches::new(len, 0, patch_indices, patch_values, None)?;
+
+        Ok(LazyPatchedArray { inner, patches })
+    }
+
+    fn with_children(array: &mut Self::Array, mut children: Vec<ArrayRef>) -> VortexResult<()> {
+        vortex_ensure_eq!(children.len(), 3);
+
+        array.inner = children.remove(0);
+
+        let patch_indices = children.remove(0);
+        let patch_values = children.remove(0);
+
+        array.patches = Patches::new(array.inner.len(), 0, patch_indices, patch_values, None)?;
+
+        Ok(())
+    }
+
+    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
+        // Execution => actually transpose the patches, get back a `PatchedArray`.
+        let patched =
+            PatchedArray::from_array_and_patches(array.inner.clone(), &array.patches, ctx)?
+                .into_array();
+
+        Ok(ExecutionResult::done(patched))
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct LazyPatchedArray {
+    inner: ArrayRef,
+    patches: Patches,
+}

vortex-array/src/arrays/lazy_patched/vtable/operations.rs

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::DynArray;
+use crate::arrays::lazy_patched::LazyPatched;
+use crate::arrays::lazy_patched::LazyPatchedArray;
+use crate::scalar::Scalar;
+use crate::vtable::OperationsVTable;
+
+impl OperationsVTable<LazyPatched> for LazyPatched {
+    fn scalar_at(array: &LazyPatchedArray, index: usize) -> VortexResult<Scalar> {
+        Ok(if let Some(scalar) = array.patches.get_patched(index)? {
+            scalar
+        } else {
+            array.inner.scalar_at(index)?
+        })
+    }
+}

vortex-array/src/arrays/lazy_patched/vtable/validity.rs

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::ArrayRef;
+use crate::arrays::lazy_patched::LazyPatched;
+use crate::arrays::lazy_patched::LazyPatchedArray;
+use crate::vtable::ValidityChild;
+
+impl ValidityChild<LazyPatched> for LazyPatched {
+    fn validity_child(array: &LazyPatchedArray) -> &ArrayRef {
+        &array.inner
+    }
+}

vortex-array/src/arrays/mod.rs

@@ -104,3 +104,4 @@ pub use variant::VariantArray;
 
 #[cfg(feature = "arbitrary")]
 pub mod arbitrary;
+pub mod lazy_patched;