Write Widen arrays through an array plugin

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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

@@ -16,8 +16,9 @@
 //! Take executes its gather instead of leaving a lazy [`Dict`] under a `Widen`, because code such
 //! as the DuckDB exporter and the dict layout writer looks for a `Dict` at the root of an array.
 //!
-//! `Widen` has no wire form yet, so the default [`ArraySession`] does not register it, and an
-//! array that holds a `Widen` cannot be serialized.
+//! `Widen` has no wire form yet. The default [`ArraySession`] registers [`WidenPlugin`], which
+//! writes each `Widen` as its canonical integer array. Serialization keeps the `Widen` statistics,
+//! and readers never see `vortex.widen`.
 //!
 //! [`ArraySession`]: crate::session::ArraySession
 //! [`Dict`]: crate::arrays::Dict
@@ -35,6 +36,7 @@ mod encode;
 mod vtable;
 use vortex_session::VortexSession;
 pub use vtable::Widen;
+pub use vtable::WidenPlugin;
 
 use crate::array::VTable;
 use crate::arrays::Dict;

vortex-array/src/arrays/widen/tests/mod.rs

@@ -10,6 +10,7 @@ mod aggregates;
 mod compute;
 mod encoding;
 mod numeric;
+mod serde;
 
 use std::sync::LazyLock;
 

vortex-array/src/arrays/widen/tests/serde.rs

@@ -0,0 +1,195 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use rstest::rstest;
+use vortex_buffer::ByteBuffer;
+use vortex_buffer::ByteBufferMut;
+use vortex_buffer::buffer;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+use vortex_session::registry::ReadContext;
+
+use super::SESSION;
+use crate::ArrayContext;
+use crate::ArrayId;
+use crate::ArrayRef;
+use crate::ArrayVTable;
+use crate::IntoArray;
+use crate::VortexSessionExecute;
+use crate::arrays::Bool;
+use crate::arrays::ChunkedArray;
+use crate::arrays::Primitive;
+use crate::arrays::PrimitiveArray;
+use crate::arrays::StructArray;
+use crate::arrays::Widen;
+use crate::arrays::WidenArray;
+use crate::arrays::widen::WidenPlugin;
+use crate::assert_arrays_eq;
+use crate::dtype::DType;
+use crate::dtype::Nullability;
+use crate::dtype::PType;
+use crate::expr::stats::Precision;
+use crate::expr::stats::Stat;
+use crate::expr::stats::StatsProvider;
+use crate::scalar::Scalar;
+use crate::serde::SerializeOptions;
+use crate::serde::SerializedArray;
+use crate::session::ArraySession;
+use crate::session::ArraySessionExt;
+use crate::validity::Validity;
+
+fn serialize_blob(
+    array: &ArrayRef,
+    ctx: &ArrayContext,
+    session: &VortexSession,
+) -> VortexResult<ByteBuffer> {
+    let mut blob = ByteBufferMut::empty();
+    for buffer in array.serialize(ctx, session, &SerializeOptions::default())? {
+        blob.extend_from_slice(buffer.as_ref());
+    }
+    Ok(blob.freeze())
+}
+
+/// Returns a `Widen` of the logical dtype `logical` over `values`, with the nullability of
+/// `values`.
+fn widen_over(values: impl IntoArray, logical: DType) -> VortexResult<ArrayRef> {
+    let values = values.into_array();
+    let logical = logical.with_nullability(values.dtype().nullability());
+    Ok(WidenArray::try_new(values, logical)?.into_array())
+}
+
+/// Serializes `array` with an unrestricted context and decodes it again. Returns the decoded
+/// array and the IDs that the context interned.
+fn serialize_round_trip(
+    array: &ArrayRef,
+    session: &VortexSession,
+) -> VortexResult<(ArrayRef, Vec<ArrayId>)> {
+    let ctx = ArrayContext::empty();
+    let serialized = SerializedArray::try_from(serialize_blob(array, &ctx, session)?)?;
+    let ids = ctx.to_ids();
+    let decoded = serialized.decode(
+        array.dtype(),
+        array.len(),
+        &ReadContext::new(ids.clone()),
+        session,
+    )?;
+
+    Ok((decoded, ids))
+}
+
+/// Asserts that `decoded` is a canonical leaf at the canonical width of its dtype.
+#[track_caller]
+fn assert_canonical_leaf(decoded: &ArrayRef) {
+    assert!(decoded.is::<Primitive>(), "{}", decoded.encoding_id());
+}
+
+#[rstest]
+#[case::integer(widen_over(buffer![-1i8, 0, 100], DType::from(PType::I64)))]
+#[case::integer_nullable(widen_over(
+    PrimitiveArray::from_option_iter([Some(-300i16), None, Some(300)]),
+    DType::from(PType::I32),
+))]
+#[case::empty(widen_over(
+    PrimitiveArray::from_iter(Vec::<i16>::new()),
+    DType::from(PType::I64),
+))]
+#[case::all_null(widen_over(
+    PrimitiveArray::new(buffer![i8::MIN, i8::MAX], Validity::AllInvalid),
+    DType::from(PType::I64),
+))]
+fn serialization_writes_widen(#[case] widen: VortexResult<ArrayRef>) -> VortexResult<()> {
+    let widen = widen?;
+    assert!(widen.is::<Widen>());
+
+    let (decoded, ids) = serialize_round_trip(&widen, &SESSION)?;
+
+    assert_canonical_leaf(&decoded);
+    assert!(!ids.contains(&ArrayVTable::id(&Widen)), "{ids:?}");
+    assert_arrays_eq!(decoded, widen, &mut SESSION.create_execution_ctx());
+
+    Ok(())
+}
+
+#[test]
+fn serialization_writes_nested_widen() -> VortexResult<()> {
+    let integer = widen_over(
+        PrimitiveArray::from_option_iter([Some(1i8), None, Some(-7)]),
+        DType::from(PType::I64),
+    )?;
+    let nested = StructArray::from_fields(&[("integer", integer.clone())])?.into_array();
+    let chunked = ChunkedArray::try_new(
+        vec![
+            integer,
+            PrimitiveArray::from_option_iter([Some(1000i64)]).into_array(),
+        ],
+        DType::Primitive(PType::I64, Nullability::Nullable),
+    )?
+    .into_array();
+
+    for array in [nested, chunked] {
+        let (decoded, ids) = serialize_round_trip(&array, &SESSION)?;
+        assert!(!ids.contains(&ArrayVTable::id(&Widen)), "{ids:?}");
+        assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx());
+    }
+
+    Ok(())
+}
+
+#[test]
+fn serialization_keeps_widen_statistics() -> VortexResult<()> {
+    let widen = widen_over(buffer![-3i8, 5, 1], DType::from(PType::I64))?;
+    let min = Scalar::primitive(-3i64, Nullability::NonNullable);
+    let max = Scalar::primitive(5i64, Nullability::NonNullable);
+
+    for (stat, value) in [(Stat::Min, &min), (Stat::Max, &max)] {
+        let value = value.value().cloned();
+        let value = value.vortex_expect("the planted statistic is not null");
+        widen.statistics().set(stat, Precision::Exact(value));
+    }
+
+    let (decoded, _) = serialize_round_trip(&widen, &SESSION)?;
+
+    assert_eq!(decoded.statistics().get(Stat::Min), Precision::Exact(min));
+    assert_eq!(decoded.statistics().get(Stat::Max), Precision::Exact(max));
+
+    Ok(())
+}
+
+#[test]
+fn serialization_writes_widen_with_restricted_context() -> VortexResult<()> {
+    let ctx = ArrayContext::new(vec![ArrayVTable::id(&Primitive)])
+        .with_allowed_ids([ArrayVTable::id(&Primitive)].into_iter().collect());
+    let widen = widen_over(buffer![-1i8, 0, 100], DType::from(PType::I64))?;
+
+    let serialized = SerializedArray::try_from(serialize_blob(&widen, &ctx, &SESSION)?)?;
+    let read_ctx = ReadContext::new(ctx.to_ids());
+    let decoded = serialized.decode(widen.dtype(), widen.len(), &read_ctx, &SESSION)?;
+
+    assert!(decoded.is::<Primitive>());
+    assert_arrays_eq!(decoded, widen, &mut SESSION.create_execution_ctx());
+
+    Ok(())
+}
+
+#[test]
+fn empty_session_can_register_widen_writer() -> VortexResult<()> {
+    let session = VortexSession::empty().with_some(ArraySession::empty());
+    session.arrays().register(Primitive);
+    session.arrays().register(Bool);
+    session.arrays().register(WidenPlugin);
+    let widen = widen_over(buffer![-1i8, 0, 100], DType::from(PType::I64))?;
+
+    let (decoded, ids) = serialize_round_trip(&widen, &session)?;
+
+    assert!(
+        !session
+            .arrays()
+            .registry()
+            .contains_key(&ArrayVTable::id(&Widen))
+    );
+    assert!(!ids.contains(&ArrayVTable::id(&Widen)));
+    assert_arrays_eq!(decoded, widen, &mut session.create_execution_ctx());
+
+    Ok(())
+}

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

@@ -2,6 +2,9 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 mod operations;
+
+mod plugin;
+pub use plugin::WidenPlugin;
 mod validity;
 
 use vortex_error::VortexResult;
@@ -91,8 +94,8 @@ impl VTable for Widen {
         _array: ArrayView<'_, Self>,
         _session: &VortexSession,
     ) -> VortexResult<Option<Vec<u8>>> {
-        // `ArraySession` does not register `Widen`, so `array_serialize` fails before it calls
-        // this method. Only a session that registers `Widen` itself can reach it.
+        // The default session registers `WidenPlugin`, which writes a storable representation
+        // without calling this method. Direct registration of the vtable has no wire form.
         vortex_panic!("TODO(connor)[Widen]: Widen has no wire form yet (#10733)")
     }
 

vortex-array/src/arrays/widen/vtable/plugin.rs

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Serialization of `Widen` arrays through existing wire forms.
+//!
+//! `Widen` has no wire form of its own, so its plugin registers only a writer.
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_session::VortexSession;
+
+use crate::ArrayDeserialization;
+use crate::ArrayId;
+use crate::ArrayPlugin;
+use crate::ArrayRef;
+use crate::ArraySerialization;
+use crate::ArrayVTable;
+use crate::IntoArray;
+use crate::RecursiveCanonical;
+use crate::VortexSessionExecute;
+use crate::arrays::Primitive;
+use crate::arrays::Widen;
+use crate::session::ArraySessionExt;
+
+/// Writes [`Widen`] arrays using an existing encoding's wire form.
+#[derive(Clone, Debug)]
+pub struct WidenPlugin;
+
+impl ArrayPlugin for WidenPlugin {
+    fn id(&self) -> ArrayId {
+        ArrayVTable::id(&Widen)
+    }
+
+    fn read_ids(&self) -> Vec<ArrayId> {
+        vec![]
+    }
+
+    fn write_ids(&self) -> Vec<ArrayId> {
+        vec![ArrayVTable::id(&Primitive)]
+    }
+
+    fn serialize(
+        &self,
+        array: &ArrayRef,
+        session: &VortexSession,
+    ) -> VortexResult<Option<ArraySerialization>> {
+        // TODO(connor)[Widen]: Add `vortex.widen` to both ID lists when it gets a wire form
+        // (#10733), and write it only where the writer permits it. This method cannot see the
+        // permitted IDs today.
+
+        // The child validity can be lazy, so the writer needs recursive canonicalization to
+        // produce nodes that it can store.
+        let canonical = array
+            .clone()
+            .execute::<RecursiveCanonical>(&mut session.create_execution_ctx())?
+            .0
+            .into_array();
+
+        session.array_serialize(&canonical)
+    }
+
+    fn deserialize(
+        &self,
+        _parts: ArrayDeserialization<'_>,
+        _session: &VortexSession,
+    ) -> VortexResult<ArrayRef> {
+        vortex_bail!("Widen has no wire form yet (#10733)")
+    }
+}

vortex-array/src/session/mod.rs

@@ -38,6 +38,7 @@ use crate::arrays::Union;
 use crate::arrays::VarBin;
 use crate::arrays::VarBinView;
 use crate::arrays::Variant;
+use crate::arrays::widen::WidenPlugin;
 
 /// Registry of array encodings.
 pub type ArrayRegistry = ArcSwapMap<Id, ArrayPluginRef>;
@@ -109,6 +110,7 @@ impl Default for ArraySession {
         this.register(PiecewiseSequence);
         this.register(RunEnd);
         this.register(VarBin);
+        this.register(WidenPlugin);
 
         this
     }

vortex-ipc/src/messages/decoder.rs

@@ -170,11 +170,21 @@ impl MessageDecoder {
 mod test {
     use bytes::BytesMut;
     use vortex_array::ArrayRef;
+    use vortex_array::ArrayVTable;
     use vortex_array::IntoArray;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::ConstantArray;
+    use vortex_array::arrays::Primitive;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::arrays::Widen;
+    use vortex_array::arrays::WidenArray;
     use vortex_array::assert_arrays_eq;
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
     use vortex_error::vortex_panic;
 
     use super::*;
@@ -220,4 +230,36 @@ mod test {
         assert_eq!(array.nbuffers(), 1, "Array should have a single buffer");
         write_and_read(&array.into_array());
     }
+
+    /// IPC writes a `Widen` as its canonical array, so the decoded encoding differs from the
+    /// written one and this test cannot use `write_and_read`.
+    #[test]
+    fn array_ipc_writes_widen() -> VortexResult<()> {
+        let values = PrimitiveArray::new(
+            buffer![-125i16, 50, 900],
+            Validity::from_iter([true, false, true]),
+        );
+        let dtype = DType::Primitive(PType::I64, Nullability::Nullable);
+        let expected = WidenArray::try_new(values.into_array(), dtype)?.into_array();
+
+        let mut ipc_bytes = BytesMut::new();
+        let mut encoder = MessageEncoder::new(SESSION.clone());
+        for buf in encoder.encode(EncoderMessage::Array(&expected))? {
+            ipc_bytes.extend_from_slice(buf.as_ref());
+        }
+
+        let mut buffer = BytesMut::from(ipc_bytes.as_ref());
+        let (array_parts, ctx, row_count) =
+            match MessageDecoder::default().read_next(&mut buffer)? {
+                PollRead::Some(DecoderMessage::Array(array_parts)) => array_parts,
+                otherwise => vortex_panic!("Expected an array, got {:?}", otherwise),
+            };
+        let actual = array_parts.decode(expected.dtype(), row_count, &ctx, &SESSION)?;
+
+        assert!(!ctx.ids().contains(&ArrayVTable::id(&Widen)));
+        assert!(actual.is::<Primitive>(), "{}", actual.encoding_id());
+        assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx());
+
+        Ok(())
+    }
 }