Benchmark normalize and fix what it found

Adds `encode_non_nullable` and `encode_nullable` arms to
`vortex-tensor/benches/normalized.rs`, so the file covers both directions of the
encoding rather than decode alone. They vary width over the same `WIDTHS` as the
decode arms, which brackets where the cost lands: a narrow vector leaves the
per-row work visible, a wide one buries it under the per-element division. The
arms reuse the file's `ELEMENTS` budget, sized for CodSpeed's CPU simulation
rather than a desktop, which is what holds every case under the 1 ms
per-iteration limit.

The benchmark contradicted the previous commit at embedding width, and found
two things.

Reading the validity off the unexecuted norms array left `L2Norm` to run a
second time under the `fill_null` execution, so the cost grew with the tensor
width. Canonicalizing first computes the norms once.

`fill_null` on a column with no nulls has nothing to fill and still charges a
cast, so it is skipped there.

Against the hand-rolled mask loop, measured on desktop-sized versions of these
arms at dimensions 8 and 768 by the fastest sample over two runs: 12-17% faster
at dimension 8, where the numbers are stable to within 0.1% across runs, and
neutral at dimension 768, where run-to-run variance of roughly 10% exceeds any
difference.

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

Author: connortsui20

vortex-tensor/benches/normalized.rs

@@ -1,13 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Baseline throughput for decoding the `Normalized` encoding over tensor columns.
+//! Baseline throughput for the `Normalized` encoding over tensor columns, in both directions:
+//! `non_nullable` and `nullable` decode, `encode_non_nullable` and `encode_nullable` split a plain
+//! tensor column into the two children.
 //!
 //! The arms vary vector width and input nullability. Their names are intended to remain stable
 //! across implementation changes so CodSpeed can compare them against `develop`.
 //!
 //! Rows are derived from a fixed element budget rather than fixed per arm, so widening a vector
 //! trades rows for elements instead of multiplying the work. See [`ELEMENTS`].
+//!
+//! Width is the axis that separates the two costs the encode path pays. A narrow vector leaves the
+//! per-row work visible, while a wide one buries it under the per-element division.
 
 #![expect(clippy::unwrap_used)]
 
@@ -19,14 +24,16 @@ use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
+use vortex_array::arrays::MaskedArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_tensor::encodings::normalized::Normalized;
+use vortex_tensor::encodings::normalized::normalize;
 use vortex_tensor::vector::Vector;
 
-// Decoding allocates the output inside the timed region, so use the vendored allocator instead
-// of measuring glibc differences between CodSpeed runner images.
+// Both directions allocate their output inside the timed region, so use the vendored allocator
+// instead of measuring glibc differences between CodSpeed runner images.
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
@@ -55,7 +62,24 @@ fn normalized_vectors(width: usize) -> ArrayRef {
     Vector::try_new_vector_array(storage).unwrap()
 }
 
-/// Masks every eighth row. The row count is `ELEMENTS / width`, matching the vector builder.
+/// Vectors of varying magnitude, so no row is a zero vector and every row takes the dividing
+/// branch in [`normalize`].
+fn plain_vectors(width: usize) -> ArrayRef {
+    let rows = ELEMENTS / width;
+    let elements: Buffer<f64> = (0..rows)
+        .flat_map(|row| (0..width).map(move |i| (row + 1) as f64 * (i + 1) as f64))
+        .collect();
+    let storage = FixedSizeListArray::new(
+        elements.into_array(),
+        u32::try_from(width).unwrap(),
+        Validity::NonNullable,
+        rows,
+    )
+    .into_array();
+    Vector::try_new_vector_array(storage).unwrap()
+}
+
+/// Masks every eighth row. The row count is `ELEMENTS / width`, matching both vector builders.
 fn sparse_nulls(width: usize) -> Validity {
     Validity::from_iter((0..ELEMENTS / width).map(|i| !i.is_multiple_of(8)))
 }
@@ -90,6 +114,15 @@ fn bench_decode(bencher: Bencher, normalized: ArrayRef, validity: Validity) {
         });
 }
 
+fn bench_encode(bencher: Bencher, input: ArrayRef) {
+    let session = vortex_array::array_session();
+    let rows = input.len();
+    bencher
+        .counter(ItemsCount::new(rows))
+        .with_inputs(|| (input.clone(), session.create_execution_ctx()))
+        .bench_values(|(input, mut ctx)| normalize(input, &mut ctx).unwrap());
+}
+
 #[divan::bench(args = WIDTHS)]
 fn non_nullable(bencher: Bencher, width: usize) {
     bench_decode(bencher, normalized_vectors(width), Validity::NonNullable);
@@ -99,3 +132,17 @@ fn non_nullable(bencher: Bencher, width: usize) {
 fn nullable(bencher: Bencher, width: usize) {
     bench_decode(bencher, normalized_vectors(width), sparse_nulls(width));
 }
+
+#[divan::bench(args = WIDTHS)]
+fn encode_non_nullable(bencher: Bencher, width: usize) {
+    bench_encode(bencher, plain_vectors(width));
+}
+
+#[divan::bench(args = WIDTHS)]
+fn encode_nullable(bencher: Bencher, width: usize) {
+    let input = MaskedArray::try_new(plain_vectors(width), sparse_nulls(width))
+        .unwrap()
+        .into_array();
+
+    bench_encode(bencher, input);
+}

vortex-tensor/src/encodings/normalized/compress.rs

@@ -159,16 +159,28 @@ pub fn normalize(input: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Normal
         .try_new_array(row_count, EmptyOptions, [input.clone()])?
         .execute(ctx)?;
 
+    // Canonicalize before reading the validity so the norms are computed once. Taking the validity
+    // off the unexecuted array instead leaves the norms to be evaluated again below, which costs
+    // the whole `L2Norm` pass a second time and grows with the tensor width.
+    let primitive_norms: PrimitiveArray = norms_array.execute(ctx)?;
+
     // `L2Norm` propagates the input's validity, so this is the column's null map.
-    let validity = norms_array.validity()?;
-
-    // Filling the nulls with zero is what moves them off the child, and it leaves the row loop
-    // below a single rule to follow: a zero norm means a zeroed row. A non-nullable input makes
-    // this a cast rather than a copy.
-    let element_dtype = DType::Primitive(tensor_match.element_ptype(), Nullability::NonNullable);
-    let norms: PrimitiveArray = norms_array
-        .fill_null(Scalar::zero_value(&element_dtype))?
-        .execute(ctx)?;
+    let validity = primitive_norms.validity()?;
+
+    // Filling the nulls with zero moves them off the child and leaves the row loop below a single
+    // rule to follow: a zero norm means a zeroed row. A column with no nulls has nothing to fill,
+    // and `fill_null` still charges it a cast, so skip it there.
+    let norms: PrimitiveArray = if validity.nullability().is_nullable() {
+        let element_dtype =
+            DType::Primitive(tensor_match.element_ptype(), Nullability::NonNullable);
+
+        primitive_norms
+            .into_array()
+            .fill_null(Scalar::zero_value(&element_dtype))?
+            .execute(ctx)?
+    } else {
+        primitive_norms
+    };
 
     let input: ExtensionArray = input.execute(ctx)?;
     let normalized_dtype = input.dtype().as_nonnullable();