Added explicit encode/decode functions to the IndexVTable

Added some tests with golden files for serializing/deserializing index
layouts.

Author: thorfour

vortex-layout/examples/reverse_index/main.rs

@@ -52,7 +52,6 @@ use vortex_array::scalar_fn::fns::operators::Operator;
 use vortex_array::search_sorted::SearchSorted;
 use vortex_array::search_sorted::SearchSortedSide;
 use vortex_array::stream::ArrayStreamExt;
-use vortex_array::stream::SendableArrayStream;
 use vortex_array::validity::Validity;
 use vortex_buffer::ByteBufferMut;
 use vortex_edition::Edition;
@@ -117,10 +116,10 @@ impl ReverseIndex {
     }
 }
 
-/// One decoded chunk of the index: its sorted keys and their serialized posting lists.
+/// A run of index entries: sorted, distinct, non-null keys and the rows holding each.
 pub struct Entries {
     keys: ArrayRef,
-    postings: VarBinViewArray,
+    postings: Vec<RoaringBitmap>,
 }
 
 impl IndexVTable for ReverseIndex {
@@ -135,6 +134,10 @@ impl IndexVTable for ReverseIndex {
         *ID
     }
 
+    fn serialize_options(&self, _options: &()) -> Vec<u8> {
+        vec![]
+    }
+
     fn deserialize_options(&self, _options: &[u8]) -> VortexResult<()> {
         Ok(())
     }
@@ -200,23 +203,57 @@ impl IndexVTable for ReverseIndex {
         }))
     }
 
+    fn encode(
+        &self,
+        chunk: &Entries,
+        _options: &(),
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef> {
+        let mut lists = Vec::with_capacity(chunk.postings.len());
+        for bitmap in &chunk.postings {
+            let mut buffer = Vec::with_capacity(bitmap.serialized_size());
+            bitmap
+                .serialize_into(&mut buffer)
+                .map_err(|err| vortex_err!("Failed to serialize postings: {err}"))?;
+            lists.push(buffer);
+        }
+
+        Ok(StructArray::try_new_with_dtype(
+            vec![
+                chunk.keys.clone(),
+                VarBinViewArray::from_iter_bin(lists).into_array(),
+            ],
+            index_fields(chunk.keys.dtype().clone()),
+            chunk.keys.len(),
+            Validity::NonNullable,
+        )?
+        .into_array())
+    }
+
     fn decode(
         &self,
         chunk: ArrayRef,
         _options: &(),
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Entries> {
         let entries = chunk.execute::<StructArray>(ctx)?;
+        let lists = entries
+            .unmasked_field_by_name(POSTINGS_FIELD)?
+            .clone()
+            .execute::<VarBinViewArray>(ctx)?;
+        let postings = (0..lists.len())
+            .map(|idx| {
+                RoaringBitmap::deserialize_from(lists.bytes_at(idx).as_slice())
+                    .map_err(|err| vortex_err!("Failed to deserialize postings: {err}"))
+            })
+            .collect::<VortexResult<_>>()?;
         Ok(Entries {
             keys: entries
                 .unmasked_field_by_name(KEY_FIELD)?
                 .clone()
                 .execute::<Canonical>(ctx)?
                 .into_array(),
-            postings: entries
-                .unmasked_field_by_name(POSTINGS_FIELD)?
-                .clone()
-                .execute::<VarBinViewArray>(ctx)?,
+            postings,
         })
     }
 
@@ -229,16 +266,13 @@ impl IndexVTable for ReverseIndex {
     ) -> VortexResult<RowLocator> {
         // Keys are unique and sorted across the partition, so at most one chunk holds the target.
         for chunk in chunks {
-            let Some(idx) = chunk
+            if let Some(idx) = chunk
                 .keys
                 .search_sorted(target, SearchSortedSide::Left)?
                 .to_found()
-            else {
-                continue;
-            };
-            let bitmap = RoaringBitmap::deserialize_from(chunk.postings.bytes_at(idx).as_slice())
-                .map_err(|err| vortex_err!("Failed to deserialize postings: {err}"))?;
-            return Ok(RowLocator::Rows(bitmap));
+            {
+                return Ok(RowLocator::Rows(chunk.postings[idx].clone()));
+            }
         }
         Ok(RowLocator::empty_rows())
     }
@@ -258,6 +292,9 @@ pub struct Builder {
 }
 
 impl IndexBuilder for Builder {
+    type Options = ();
+    type Chunk = Entries;
+
     fn push(
         &mut self,
         chunk: &ArrayRef,
@@ -277,12 +314,12 @@ impl IndexBuilder for Builder {
         Ok(())
     }
 
-    fn finish(self: Box<Self>) -> VortexResult<Option<(SendableArrayStream, Vec<u8>)>> {
+    fn finish(self) -> VortexResult<Option<(Vec<Entries>, ())>> {
         let Builder {
             dtype,
             postings,
             allocator,
-        } = *self;
+        } = self;
 
         let mut entries: Vec<(Scalar, RoaringBitmap)> = postings.into_iter().collect();
         entries.sort_by(|(a, _), (b, _)| {
@@ -291,31 +328,22 @@ impl IndexBuilder for Builder {
         });
 
         let key_dtype = dtype.as_nonnullable();
-        let mut key_builder = builder_with_capacity_in(&key_dtype, entries.len(), &allocator);
-        let mut lists = Vec::with_capacity(entries.len());
-        for (key, bitmap) in &entries {
+        let mut keys = builder_with_capacity_in(&key_dtype, entries.len(), &allocator);
+        let mut bitmaps = Vec::with_capacity(entries.len());
+        for (key, bitmap) in entries {
             // Keys are never null (`push` skips them), but may carry the source column's nullable
             // dtype; the key column itself is non-nullable, so normalize before appending.
-            key_builder.append_scalar(&key.cast(&key_dtype)?)?;
-            let mut buffer = Vec::with_capacity(bitmap.serialized_size());
-            bitmap
-                .serialize_into(&mut buffer)
-                .map_err(|err| vortex_err!("Failed to serialize postings: {err}"))?;
-            lists.push(buffer);
+            keys.append_scalar(&key.cast(&key_dtype)?)?;
+            bitmaps.push(bitmap);
         }
 
-        let len = entries.len();
-        let array = StructArray::try_new_with_dtype(
-            vec![
-                key_builder.finish(),
-                VarBinViewArray::from_iter_bin(lists).into_array(),
-            ],
-            index_fields(key_dtype),
-            len,
-            Validity::NonNullable,
-        )?;
-
-        Ok(Some((array.into_array().to_array_stream().boxed(), vec![])))
+        Ok(Some((
+            vec![Entries {
+                keys: keys.finish(),
+                postings: bitmaps,
+            }],
+            (),
+        )))
     }
 
     fn buffered_bytes(&self) -> u64 {

vortex-layout/examples/reverse_index/tests.rs

@@ -1,17 +1,26 @@
 //! End-to-end tests: write a column indexed by [`ReverseIndex`], then probe it.
 
+use std::io::Write;
 use std::sync::Arc;
 
+use goldenfile::Mint;
+use roaring::RoaringBitmap;
+use rstest::rstest;
+
 use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
 use vortex_array::MaskFuture;
 use vortex_array::VortexSessionExecute;
 use vortex_array::array_session;
+use vortex_array::assert_arrays_eq;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::arrays::struct_::StructArrayExt;
+use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldPath;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
 use vortex_array::expr::col;
 use vortex_array::expr::eq;
 use vortex_array::expr::lit;
@@ -39,13 +48,15 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 use vortex_layout::layouts::indexed::INDEXED_LAYOUT_ID;
 use vortex_layout::layouts::indexed::IndexConfig;
 use vortex_layout::layouts::indexed::IndexSessionExt;
+use vortex_layout::layouts::indexed::IndexVTable;
 use vortex_layout::layouts::indexed::IndexedStrategy;
 use vortex_layout::layouts::repartition::RepartitionStrategy;
 use vortex_layout::layouts::repartition::RepartitionWriterOptions;
 use vortex_layout::session::LayoutSession;
 use vortex_mask::Mask;
 use vortex_session::VortexSession;
 
+use crate::Entries;
 use crate::ReverseIndex;
 
 /// Small enough that a 12-row file spans three blocks, making the row/block granularity
@@ -348,3 +359,55 @@ async fn string_valued_column_is_indexed_exactly() -> VortexResult<()> {
     );
     Ok(())
 }
+
+/// `ReverseIndex`'s content is a contract with every file already written: it must encode as the
+/// declared index dtype, and re-encoding its decoding must reproduce that encoding.
+///
+/// This is `vortex_layout::layouts::indexed::test_harness::check_roundtrip`, which kinds in other
+/// crates can call by enabling `vortex-layout`'s `_test-harness` feature; an example in the same
+/// package cannot enable it for itself. The goldenfile pins the encoded rows themselves, so a
+/// change to how keys or postings are laid out shows up as a diff against the checked-in file.
+#[rstest]
+#[case::i32(
+    "reverse_index_i32.txt",
+    PrimitiveArray::from_iter([10i32, 20, 30]).into_array(),
+    DType::Primitive(PType::I32, Nullability::Nullable),
+)]
+#[case::utf8(
+    "reverse_index_utf8.txt",
+    VarBinViewArray::from_iter_str(["alpha", "beta", "gamma"]).into_array(),
+    DType::Utf8(Nullability::NonNullable),
+)]
+fn reverse_index_roundtrips(
+    #[case] golden: &str,
+    #[case] keys: ArrayRef,
+    #[case] data_dtype: DType,
+) -> VortexResult<()> {
+    let mut ctx = session().create_execution_ctx();
+    let entries = Entries {
+        keys,
+        postings: vec![
+            RoaringBitmap::from_iter([0u32, 9]),
+            RoaringBitmap::new(),
+            RoaringBitmap::from_iter([3u32, 4, 100_000]),
+        ],
+    };
+
+    let encoded = ReverseIndex.encode(&entries, &(), &mut ctx)?;
+    assert_eq!(
+        Some(encoded.dtype()),
+        ReverseIndex.index_dtype(&data_dtype, &()).as_ref()
+    );
+    let decoded = ReverseIndex.decode(encoded.clone(), &(), &mut ctx)?;
+    assert_arrays_eq!(
+        ReverseIndex.encode(&decoded, &(), &mut ctx)?,
+        encoded,
+        &mut ctx
+    );
+
+    let mut mint = Mint::new("goldenfiles/");
+    let mut file = mint.new_goldenfile(golden)?;
+    writeln!(file, "{}", encoded.display_values())?;
+    Ok(())
+}
+

vortex-layout/goldenfiles/indexed.metadata

@@ -0,0 +1 @@
+[{key: 10i32, postings: "3a 30 0 0 1 0 0 0 0 0 1 0 10 0 0 0 0 0 9 0"}, {key: 20i32, postings: "3a 30 0 0 0 0 0 0"}, {key: 30i32, postings: "3a 30 0 0 2 0 0 0 0 0 1 0 1 0 0 0 18 0 0 0 1c 0 0 0 3 0 4 0 a0 86"}]

vortex-layout/goldenfiles/reverse_index_i32.txt

@@ -0,0 +1 @@
+[{key: "alpha", postings: "3a 30 0 0 1 0 0 0 0 0 1 0 10 0 0 0 0 0 9 0"}, {key: "beta", postings: "3a 30 0 0 0 0 0 0"}, {key: "gamma", postings: "3a 30 0 0 2 0 0 0 0 0 1 0 1 0 0 0 18 0 0 0 1c 0 0 0 3 0 4 0 a0 86"}]

vortex-layout/goldenfiles/reverse_index_utf8.txt

@@ -11,7 +11,6 @@ use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::dtype::DType;
 use vortex_array::expr::BoundExpression;
-use vortex_array::stream::SendableArrayStream;
 use vortex_buffer::BitBufferMut;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
@@ -39,11 +38,25 @@ pub type IndexId = Id;
 /// [`Chunk`](Self::Chunk) at most once, sharing it between every expression that probes the
 /// chunk. Each expression is reduced by [`plan`](Self::plan) to a [`Query`](Self::Query) that
 /// [`resolve`](Self::resolve) answers against those decoded chunks.
+///
+/// # Serialization
+///
+/// A kind's on-disk format is exactly three things, and nothing else a kind does reaches the file:
+///
+/// - [`serialize_options`](Self::serialize_options) and
+///   [`deserialize_options`](Self::deserialize_options), for the options blob stored in the spec.
+///   The blob is the kind's to version.
+/// - [`index_dtype`](Self::index_dtype), the schema of the index child.
+/// - [`encode`](Self::encode) and [`decode`](Self::decode), between a [`Chunk`](Self::Chunk) and
+///   rows of that schema.
+///
+/// Builders produce chunks and options, never bytes, so these pairs can be tested without
+/// writing a layout. Re-encoding a decoded chunk must reproduce its encoding.
 pub trait IndexVTable: 'static + Send + Sync + Debug {
     /// The kind's options, parsed from the self-versioned options blob.
     type Options: 'static + Send + Sync;
     /// Builds the index on the write path.
-    type Builder: IndexBuilder + 'static;
+    type Builder: IndexBuilder<Options = Self::Options, Chunk = Self::Chunk> + 'static;
     /// What [`plan`](Self::plan) reduced one expression to, answered by
     /// [`resolve`](Self::resolve).
     type Query: 'static + Send + Sync;
@@ -53,6 +66,9 @@ pub trait IndexVTable: 'static + Send + Sync + Debug {
     /// Stable string id, e.g. `vortex.idx.reverse_index`.
     fn id(&self) -> IndexId;
 
+    /// Serialize options into the blob stored in the spec.
+    fn serialize_options(&self, options: &Self::Options) -> Vec<u8>;
+
     /// Parse an options blob, as configured on the write side or stored in a spec.
     fn deserialize_options(&self, options: &[u8]) -> VortexResult<Self::Options>;
 
@@ -94,9 +110,20 @@ pub trait IndexVTable: 'static + Send + Sync + Debug {
         options: &Self::Options,
     ) -> VortexResult<Option<IndexQueryPlan<Self::Query>>>;
 
-    /// Decode one chunk of the index child, in the index child's own schema.
+    /// Encode a chunk a builder produced as rows of the index child, of the declared
+    /// [`index_dtype`](Self::index_dtype).
+    fn encode(
+        &self,
+        chunk: &Self::Chunk,
+        options: &Self::Options,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<ArrayRef>;
+
+    /// Decode rows of the index child, in its own schema, back into a chunk.
     ///
-    /// Called at most once per chunk per reader, however many expressions probe it.
+    /// The index child's layout strategy may chunk it differently than it was encoded, so this
+    /// must accept any contiguous run of encoded rows. Called at most once per chunk per reader,
+    /// however many expressions probe it.
     fn decode(
         &self,
         chunk: ArrayRef,
@@ -190,6 +217,11 @@ impl<V: IndexVTable> DynIndexVTable for V {
 
 /// Accumulates index content while the data stream is written.
 pub trait IndexBuilder: Send {
+    /// The kind's options, as finally chosen by the build.
+    type Options;
+    /// The content the kind's [`IndexVTable::encode`] writes out.
+    type Chunk;
+
     /// Chunks arrive in stream order with their row offset within this builder's partition, which
     /// for an unpartitioned index is the whole layout.
     fn push(
@@ -199,17 +231,18 @@ pub trait IndexBuilder: Send {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<()>;
 
-    /// Emit the index content as an array stream, to be written through a child layout strategy.
+    /// Emit the index content as chunks, which the writer encodes and writes through a child
+    /// layout strategy in order.
     ///
-    /// Returns the final serialized options alongside it, so builders can record normalization
-    /// choices or block sizes discovered during the build.
+    /// Returns the final options alongside them, so builders can record normalization choices or
+    /// block sizes discovered during the build.
     ///
     /// `None` declines: nothing worth keeping was built, so no index child and no spec are written,
     /// and the wrapper collapses to the plain data layout if every builder declines. This is the
     /// only point at which size can be judged — a builder is constructed before the first chunk
     /// arrives, so row count and cardinality are not knowable earlier. Declining is always safe: an
     /// absent index reads exactly like an unregistered one.
-    fn finish(self: Box<Self>) -> VortexResult<Option<(SendableArrayStream, Vec<u8>)>>;
+    fn finish(self) -> VortexResult<Option<(Vec<Self::Chunk>, Self::Options)>>;
 
     /// Bytes currently buffered, reported up through the write context's buffered-bytes tracker.
     fn buffered_bytes(&self) -> u64;